Injecting arbitrary code into PE Files using AsmResolver - https://washi.dev/blog/posts/import-patching/

Program.cs

using AsmResolver;
using AsmResolver.PE;
using AsmResolver.PE.Code;
using AsmResolver.PE.File;
using AsmResolver.PE.File.Headers;
using AsmResolver.PE.Imports;
using AsmResolver.PE.Imports.Builder;
using AsmResolver.PE.Platforms;
using AsmResolver.PE.Relocations;
using AsmResolver.PE.Relocations.Builder;

namespace ImportPatcher;

internal static class Program
{
    public static void Main(string[] args)
    {
        string path = args[0];
        
        var file = PEFile.FromFile(path);
        var image = PEImage.FromFile(file);

        // Inject a new entry point.
        var newCodeSegment = InjectNewCode(file, image);
        file.UpdateHeaders();
        file.OptionalHeader.AddressOfEntryPoint = newCodeSegment.Rva;
        
        // Inject IAT trampolines and the new import directories and accompanied relocation directories.
        InjectTrampolines(file, image);
        RebuildIatAndRelocs(file, image);

        // Write to output.
        file.Write(Path.ChangeExtension(path, $".patched{Path.GetExtension(path)}"));
    }

    private static ISegment InjectNewCode(PEFile file, IPEImage image)
    {
        // Import ucrtbase.dll!puts into the PE image.
        var ucrtbase = new ImportedModule("ucrtbase.dll");
        var puts = new ImportedSymbol(0, "puts");
        ucrtbase.Symbols.Add(puts);
        image.Imports.Add(ucrtbase);

        // Define an entry point symbol that we can reference later in our code.
        var originalEntryPoint = new Symbol(file.GetReferenceToRva(file.OptionalHeader.AddressOfEntryPoint));

        // Write some code.
        var code = image.MachineType switch
        {
            MachineType.I386 => CreateI386CodeSegment(puts, originalEntryPoint),
            MachineType.Amd64 => CreateAmd64CodeSegment(puts, originalEntryPoint),
            _ => throw new NotImplementedException($"Platform {image.MachineType} is not implemented.")
        };

        // Add all relocatable data entries to the PE.
        foreach (var relocation in code.Relocations)
            image.Relocations.Add(relocation);

        // Add actual code to a new RX section.
        file.Sections.Add(new PESection(
            ".text2",
            SectionFlags.MemoryRead | SectionFlags.MemoryExecute | SectionFlags.ContentCode,
            code.Segment));
        
        return code.Segment;
    }

    private static RelocatableSegment CreateI386CodeSegment(ISymbol puts, ISymbol originalEntryPoint)
    {
        var code = new DataSegment(new byte[]
            {
                /* 00000000: */ 0x68, 0x00, 0x00, 0x00, 0x00,       // push &message
                /* 00000005: */ 0xFF, 0x15, 0x00, 0x00, 0x00, 0x00, // call [&puts]
                /* 0000000B: */ 0xB8, 0x00, 0x00, 0x00, 0x01,       // mov eax, &originalEntryPoint
                /* 00000010: */ 0xFF, 0xE0,                         // jmp eax
                
                // message:
                /* 00000012: */ 0x48, 0x69, 0x20, 0x49, 0x20, 0x61, 0x6d, 0x20, 0x69, 0x6e, // "Hi I am in"
                /* 0000001c: */ 0x6a, 0x65, 0x63, 0x74, 0x65, 0x64, 0x20, 0x63, 0x6f, 0x64, // "jected cod"
                /* 00000026: */ 0x65, 0x21, 0x00                                            // "e!."
            }).AsPatchedSegment()
            .Patch(relativeOffset: 0x1, AddressFixupType.Absolute32BitAddress, +0x12 /* &message */) 
            .Patch(relativeOffset: 0x7, AddressFixupType.Absolute32BitAddress, puts)
            .Patch(relativeOffset: 0xC, AddressFixupType.Absolute32BitAddress, originalEntryPoint);

        return new RelocatableSegment(code, new[]
        {
            new BaseRelocation(RelocationType.HighLow, code.ToReference(0x1)),
            new BaseRelocation(RelocationType.HighLow, code.ToReference(0x7)),
            new BaseRelocation(RelocationType.HighLow, code.ToReference(0xC)),
        });
    }
    
    private static RelocatableSegment CreateAmd64CodeSegment(ISymbol puts, ISymbol originalEntryPoint)
    {
        var code = new DataSegment(new byte[]
            {
                /* 00000000: */ 0x48, 0x8D, 0x0D, 0x00, 0x00, 0x00, 0x00,                   // lea rcx, [rel message]
                /* 00000007: */ 0xFF, 0x15, 0x00, 0x00, 0x00, 0x00,                         // call [rel puts]
                /* 0000000D: */ 0x48, 0xB8, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, // mov rax, &originalEntryPoint
                /* 00000017: */ 0xFF, 0xE0,                                                 // jmp rax
                
                // message:
                /* 00000019: */ 0x48, 0x69, 0x20, 0x49, 0x20, 0x61, 0x6d, 0x20, 0x69, 0x6e, // "Hi I am in"
                /* 00000023: */ 0x6a, 0x65, 0x63, 0x74, 0x65, 0x64, 0x20, 0x63, 0x6f, 0x64, // "jected cod"
                /* 0000002d: */ 0x65, 0x21, 0x00                                            // "e!."
            }).AsPatchedSegment()
            .Patch(0x3, AddressFixupType.Relative32BitAddress, 0x19)
            .Patch(0x9, AddressFixupType.Relative32BitAddress, puts)
            .Patch(0xF, AddressFixupType.Absolute64BitAddress, originalEntryPoint);

        return new RelocatableSegment(code, new[]
        {
            new BaseRelocation(RelocationType.Dir64, code.ToReference(0xF))
        });
    }

    private static void InjectTrampolines(PEFile file, IPEImage image)
    {
        var platform = Platform.Get(file.FileHeader.Machine);
        bool is64Bit = file.OptionalHeader.Magic == OptionalHeaderMagic.PE64;
        
        // Construct the trampoline table.
        var trampolineTable = new SegmentBuilder();
        foreach (var module in image.Imports)
        {
            foreach (var originalSymbol in module.Symbols)
            {
                // Look up IAT RVA and get accompanying section. 
                if (originalSymbol.AddressTableEntry is null)
                    continue;

                uint iatEntryRva = originalSymbol.AddressTableEntry!.Rva;
                var section = file.GetSectionContainingRva(iatEntryRva);
                if (section.Contents?.AsPatchedSegment() is not { } sectionContents)
                    continue;

                // Construct trampoline stub and make it relocatable.
                var trampoline = platform.CreateThunkStub(originalSymbol);
                var trampolineSymbol = new Symbol(trampoline.Segment.ToReference());
                
                trampolineTable.Add(trampoline.Segment);
                foreach (var relocation in trampoline.Relocations)
                    image.Relocations.Add(relocation);

                // Patch original IAT entry with address to the trampoline stub, and make it relocatable.
                sectionContents.Patch(
                    iatEntryRva - sectionContents.Rva,
                    AddressFixupType.Absolute64BitAddress,
                    trampolineSymbol);

                image.Relocations.Add(new BaseRelocation(
                    is64Bit ? RelocationType.Dir64 : RelocationType.HighLow,
                    originalSymbol.AddressTableEntry));

                // Update section contents.
                section.Contents = sectionContents;
            }
        }

        // Add trampoline stubs to a new RX section.
        file.Sections.Add(new PESection(
            ".trmpln",
            SectionFlags.ContentCode | SectionFlags.MemoryRead | SectionFlags.MemoryExecute,
            trampolineTable));
    }

    private static void RebuildIatAndRelocs(PEFile file, IPEImage image)
    {
        // Rebuild import directories.
        var importBuffer = new ImportDirectoryBuffer(image.MachineType != MachineType.Amd64);
        foreach (var module in image.Imports)
            importBuffer.AddModule(module);

        file.Sections.Add(new PESection(
            ".idata2",
            SectionFlags.ContentCode | SectionFlags.ContentUninitializedData | SectionFlags.MemoryRead,
            new SegmentBuilder
            {
                {importBuffer, 8},
                {importBuffer.ImportAddressDirectory, 8},
            }));

        // Rebuild relocations directory.
        var relocationBuffer = new RelocationsDirectoryBuffer();
        foreach (var relocation in image.Relocations)
            relocationBuffer.Add(relocation);
        
        file.Sections.Add(new PESection(
            ".reloc2",
            SectionFlags.ContentInitializedData | SectionFlags.MemoryRead | SectionFlags.MemoryDiscardable,
            relocationBuffer
        ));

        // Recalculate offsets and update headers, rewiring the old IDT, IAT and reloc tables to the new ones.
        file.UpdateHeaders();
        
        var dataDirectories = file.OptionalHeader.DataDirectories;
        dataDirectories[(int) DataDirectoryIndex.BaseRelocationDirectory] = new(
            relocationBuffer.Rva,
            relocationBuffer.GetVirtualSize());
        
        dataDirectories[(int) DataDirectoryIndex.ImportDirectory] = new(
            importBuffer.Rva
            , importBuffer.GetVirtualSize());
        
        dataDirectories[(int) DataDirectoryIndex.IatDirectory] = new(
            importBuffer.ImportAddressDirectory.Rva, 
            importBuffer.ImportAddressDirectory.GetVirtualSize());
    }
}