A simple utility for modifying/adding exports to a PE file

expmod.cpp

#include <windows.h>
#include <iostream>
#include <fstream>
#include <string>
#include <vector>
#include <ctime>
#include <memory>
#include <optional>
#include <random>
#include <string_view>
#include <unordered_map>
#include <charconv>

#define DEBUG_LOGGING 0

#if DEBUG_LOGGING
#define log_debug(msg) std::cout << "[DEBUG] " << msg << std::endl
#else
#define log_debug(msg)
#endif

#define log_info(msg) std::cout << "[INFO] " << msg << std::endl
#define log_error(msg) std::cerr << "[ERROR] " << msg << std::endl

struct exp_entry
{
    std::string name;
    DWORD rva;
};

struct expmod
{
    struct mmf
    {
        HANDLE map = nullptr;
        void* data = nullptr;
        SIZE_T size = 0;

        ~mmf()
        {
            if ( data )
            {
                FlushViewOfFile( data, 0 );
                UnmapViewOfFile( data );
            }
            if ( map )
            {
                CloseHandle( map );
            }
        }
    };

    std::unique_ptr<mmf> mem;
    std::string path;

    IMAGE_DOS_HEADER* dos = nullptr;
    IMAGE_NT_HEADERS64* nt = nullptr;
    IMAGE_SECTION_HEADER* sec = nullptr;

    IMAGE_EXPORT_DIRECTORY* exp_dir = nullptr;
    DWORD exp_rva = 0;
    DWORD exp_size = 0;

    DWORD* func_addr = nullptr;
    DWORD* name_addr = nullptr;
    WORD* name_ord = nullptr;

    std::vector<exp_entry> old_exp;
    std::vector<exp_entry> new_exp;

    bool load_file( const std::string& p )
    {
        log_debug( "loading file: " << p );

        path = p;
        mem = std::make_unique<mmf>();

        HANDLE file = CreateFileA( p.c_str(), GENERIC_READ | GENERIC_WRITE,
                                   FILE_SHARE_READ, nullptr, OPEN_EXISTING,
                                   FILE_ATTRIBUTE_NORMAL, nullptr );

        if ( file == INVALID_HANDLE_VALUE )
        {
            DWORD err = GetLastError();
            log_error( "failed to open file: " << p << ", Error: " << err );
            return false;
        }

        DWORD size_high;
        DWORD size_low = GetFileSize( file, &size_high );
        mem->size = ( ( SIZE_T ) size_high << 32 ) | size_low;

        log_debug( "file size: " << mem->size << " bytes" );

        mem->map = CreateFileMappingA( file, nullptr, PAGE_READWRITE,
                                       size_high, size_low, nullptr );
        CloseHandle( file );

        if ( !mem->map )
        {
            DWORD err = GetLastError();
            log_error( "failed to create file mapping, Error: " << err );
            return false;
        }

        log_debug( "file mapping created successfully" );

        mem->data = MapViewOfFile( mem->map, FILE_MAP_ALL_ACCESS, 0, 0, 0 );

        if ( !mem->data )
        {
            DWORD err = GetLastError();
            log_error( "failed to map view of file, Error: " << err );
            return false;
        }

        log_debug( "file mapped into memory at address: " << std::hex << mem->data << std::dec );

        return init_headers();
    }

    bool init_headers()
    {
        log_debug( "initializing PE headers" );

        if ( !mem->data )
        {
            return false;
        }

        dos = static_cast< IMAGE_DOS_HEADER* >( mem->data );
        if ( dos->e_magic != IMAGE_DOS_SIGNATURE )
        {
            log_error( "invalid DOS signature" );
            return false;
        }

        nt = reinterpret_cast< IMAGE_NT_HEADERS64* >( ( BYTE* ) ( mem->data ) + dos->e_lfanew );
        if ( nt->Signature != IMAGE_NT_SIGNATURE )
        {
            log_error( "invalid NT signature" );
            return false;
        }

        if ( nt->FileHeader.Machine != IMAGE_FILE_MACHINE_AMD64 )
        {
            log_error( "not a 64b PE" );
            return false;
        }

        log_debug( "number of sections: " << nt->FileHeader.NumberOfSections );
        log_debug( "size of optional header: " << nt->FileHeader.SizeOfOptionalHeader );

        sec = reinterpret_cast< IMAGE_SECTION_HEADER* >(
            reinterpret_cast< BYTE* >( &nt->OptionalHeader ) +
            nt->FileHeader.SizeOfOptionalHeader );

        log_debug( "section headers at: " << std::hex << sec << std::dec );
        log_debug( "checking export directory" );
        log_debug( "export directory RVA: 0x" << std::hex
                   << nt->OptionalHeader.DataDirectory[ IMAGE_DIRECTORY_ENTRY_EXPORT ].VirtualAddress
                   << std::dec );
        log_debug( "export directory Size: 0x" << std::hex
                   << nt->OptionalHeader.DataDirectory[ IMAGE_DIRECTORY_ENTRY_EXPORT ].Size
                   << std::dec );

        if ( nt->OptionalHeader.DataDirectory[ IMAGE_DIRECTORY_ENTRY_EXPORT ].VirtualAddress == 0 )
        {
            log_info( "no export directory found. Will create one." );
            return true;
        }

        exp_rva = nt->OptionalHeader.DataDirectory[ IMAGE_DIRECTORY_ENTRY_EXPORT ].VirtualAddress;
        exp_size = nt->OptionalHeader.DataDirectory[ IMAGE_DIRECTORY_ENTRY_EXPORT ].Size;

        exp_dir = ( IMAGE_EXPORT_DIRECTORY* ) ( rva_to_va( exp_rva ) );

        log_debug( "export directory at VA: " << std::hex << exp_dir << std::dec );

        if ( !exp_dir )
        {
            log_error( "failed to get export directory" );
            return false;
        }

        if ( DEBUG_LOGGING )
        {
            log_debug( "export directory details:" );
            log_debug( "  Characteristics: 0x" << std::hex << exp_dir->Characteristics << std::dec );
            log_debug( "  TimeDateStamp: 0x" << std::hex << exp_dir->TimeDateStamp << std::dec );
            log_debug( "  MajorVersion: " << exp_dir->MajorVersion );
            log_debug( "  MinorVersion: " << exp_dir->MinorVersion );
            log_debug( "  Name RVA: 0x" << std::hex << exp_dir->Name << std::dec );

            char* name = ( char* ) ( rva_to_va( exp_dir->Name ) );
            if ( name )
            {
                log_debug( "  Name: " << name );
            }
            else
            {
                log_debug( "  Name: <invalid>" );
            }

            log_debug( "  Base: " << exp_dir->Base );
            log_debug( "  NumberOfFunctions: " << exp_dir->NumberOfFunctions );
            log_debug( "  NumberOfNames: " << exp_dir->NumberOfNames );
            log_debug( "  AddressOfFunctions RVA: 0x" << std::hex << exp_dir->AddressOfFunctions << std::dec );
            log_debug( "  AddressOfNames RVA: 0x" << std::hex << exp_dir->AddressOfNames << std::dec );
            log_debug( "  AddressOfNameOrdinals RVA: 0x" << std::hex << exp_dir->AddressOfNameOrdinals << std::dec );
        }

        func_addr = ( DWORD* ) ( rva_to_va( exp_dir->AddressOfFunctions ) );
        name_addr = ( DWORD* ) ( rva_to_va( exp_dir->AddressOfNames ) );
        name_ord = ( WORD* ) ( rva_to_va( exp_dir->AddressOfNameOrdinals ) );

        log_debug( "AddressOfFunctions VA: " << std::hex << func_addr << std::dec );
        log_debug( "AddressOfNames VA: " << std::hex << name_addr << std::dec );
        log_debug( "AddressOfNameOrdinals VA: " << std::hex << name_ord << std::dec );

        if ( !func_addr || !name_addr || !name_ord )
        {
            log_error( "failed to get export tables" );
            return false;
        }

        load_exports();

        return true;
    }

    void load_exports()
    {
        log_debug( "loading existing exports" );

        if ( !exp_dir )
        {
            log_debug( "no export directory, nothing to load" );
            return;
        }

        old_exp.clear();

        for ( DWORD i = 0; i < exp_dir->NumberOfNames; i++ )
        {
            char* name = ( char* ) ( rva_to_va( name_addr[ i ] ) );
            if ( !name )
            {
                log_debug( "export " << i << ": Name pointer is invalid" );
                continue;
            }

            WORD ord = name_ord[ i ];
            if ( ord >= exp_dir->NumberOfFunctions )
            {
                log_debug( "export " << i << ": ordinal " << ord << " is out of range" );
                continue;
            }

            DWORD rva = func_addr[ ord ];

            log_debug( "export " << i << ": name=" << name << ", ordinal=" << ord << ", RVA=0x" << std::hex << rva << std::dec );

            exp_entry e{ std::string( name ), rva };
            old_exp.push_back( e );
        }

        log_debug( "loaded " << old_exp.size() << " existing exports" );
    }

    void* rva_to_va( DWORD rva ) const
    {
        if ( !mem->data || rva == 0 )
        {
            return nullptr;
        }

        DWORD off = rva_to_off( rva );
        if ( off == 0 && rva != 0 )
        {
            return nullptr;
        }

        if ( off >= mem->size )
        {
            return nullptr;
        }

        void* va = ( BYTE* ) ( mem->data ) + off;
        return va;
    }

    DWORD rva_to_off( DWORD rva ) const
    {
        if ( !mem->data || rva == 0 )
        {
            return 0;
        }

        for ( int i = 0; i < nt->FileHeader.NumberOfSections; i++ )
        {
            if ( rva >= sec[ i ].VirtualAddress &&
                 rva < sec[ i ].VirtualAddress + sec[ i ].SizeOfRawData )
            {

                DWORD off = rva - sec[ i ].VirtualAddress + sec[ i ].PointerToRawData;
                return off;
            }
        }

        return rva;
    }

    DWORD off_to_rva( DWORD off ) const
    {
        if ( !mem->data || off == 0 )
        {
            return 0;
        }

        for ( int i = 0; i < nt->FileHeader.NumberOfSections; i++ )
        {
            if ( off >= sec[ i ].PointerToRawData &&
                 off < sec[ i ].PointerToRawData + sec[ i ].SizeOfRawData )
            {
                DWORD rva = off - sec[ i ].PointerToRawData + sec[ i ].VirtualAddress;
                return rva;
            }
        }

        return off;
    }

    bool add_export( const std::string& name, DWORD64 va )
    {
        log_debug( "adding export: " << name << " -> 0x" << std::hex << va << std::dec );

        if ( !mem->data )
        {
            log_error( "file data is NULL" );
            return false;
        }

        DWORD rva = static_cast< DWORD >( va - nt->OptionalHeader.ImageBase );
        log_debug( "converted VA 0x" << std::hex << va << " to RVA 0x" << rva << std::dec );

        exp_entry e{ name, rva };
        new_exp.push_back( e );

        log_debug( "added export to list, total new exports: " << new_exp.size() );

        return true;
    }

    bool update_exports()
    {
        log_debug( "adding to existing export directory" );

        if ( !mem->data )
        {
            return false;
        }

        if ( new_exp.empty() )
        {
            log_debug( "no exports to add" );
            return true;
        }

        if ( !exp_dir )
        {
            log_debug( "no existing export directory, will create new one" );
            return create_exp_dir();
        }

        log_debug( "found existing export directory" );

        int sec_idx = -1;
        for ( int i = 0; i < nt->FileHeader.NumberOfSections; i++ )
        {
            if ( exp_rva >= sec[ i ].VirtualAddress &&
                 exp_rva < sec[ i ].VirtualAddress + sec[ i ].Misc.VirtualSize )
            {
                sec_idx = i;
                break;
            }
        }

        if ( sec_idx == -1 )
        {
            log_error( "could not find section containing export directory" );
            return false;
        }

        log_debug( "export directory is in section " << sec_idx );

        char name[ 9 ] = { 0 };
        memcpy( name, sec[ sec_idx ].Name, 8 );
        log_debug( "export section name: " << name );

        std::vector<exp_entry> back = old_exp;

        if ( create_exp_dir() )
        {
            log_debug( "successfully created new export directory with all exports" );
            return true;
        }

        log_error( "failed to create new export directory" );
        return false;
    }

    bool create_new_sec( DWORD need_size, const std::vector<exp_entry>& all_exp )
    {
        log_debug( "creating new section for export directory" );

        if ( nt->FileHeader.NumberOfSections >= 96 )
        {
            log_error( "maximum number of sections reached" );
            return false;
        }

        DWORD align_size = ( need_size + 0xFFF ) & ~0xFFF;
        log_debug( "aligned section size: 0x" << std::hex << align_size << std::dec );

        IMAGE_SECTION_HEADER* last = &sec[ nt->FileHeader.NumberOfSections - 1 ];

        char lname[ 9 ] = { 0 };
        memcpy( lname, last->Name, 8 );
        log_debug( "prev section: " << lname );
        log_debug( "prev section VirtualAddress: 0x" << std::hex << last->VirtualAddress << std::dec );
        log_debug( "prev section VirtualSize: 0x" << std::hex << last->Misc.VirtualSize << std::dec );
        log_debug( "prev section PointerToRawData: 0x" << std::hex << last->PointerToRawData << std::dec );
        log_debug( "prev section SizeOfRawData: 0x" << std::hex << last->SizeOfRawData << std::dec );

        DWORD new_rva = ( last->VirtualAddress + last->Misc.VirtualSize + 0xFFF ) & ~0xFFF;
        log_debug( "new section RVA: 0x" << std::hex << new_rva << std::dec );

        DWORD align = nt->OptionalHeader.FileAlignment;
        DWORD new_off = ( last->PointerToRawData + last->SizeOfRawData + align - 1 ) & ~( align - 1 );
        log_debug( "new section raw offset: 0x" << std::hex << new_off << std::dec );

        DWORD new_size = new_off + align_size;
        log_debug( "current file size: 0x" << std::hex << mem->size << std::dec );
        log_debug( "new file size: 0x" << std::hex << new_size << std::dec );

        if ( new_size > mem->size )
        {
            log_debug( "need to expand file from 0x" << std::hex << mem->size << " to 0x" << new_size << std::dec );

            UnmapViewOfFile( mem->data );
            CloseHandle( mem->map );

            HANDLE file = CreateFileA( path.c_str(), GENERIC_READ | GENERIC_WRITE,
                                       FILE_SHARE_READ, nullptr, OPEN_EXISTING,
                                       FILE_ATTRIBUTE_NORMAL, nullptr );

            if ( file == INVALID_HANDLE_VALUE )
            {
                DWORD err = GetLastError();
                log_error( "failed to reopen file for resizing: " << path << ", Error: " << err );
                return false;
            }

            LARGE_INTEGER dist;
            dist.QuadPart = new_size;
            if ( !SetFilePointerEx( file, dist, nullptr, FILE_BEGIN ) )
            {
                DWORD err = GetLastError();
                log_error( "Failed to set file pointer: " << err );
                CloseHandle( file );
                return false;
            }

            if ( !SetEndOfFile( file ) )
            {
                DWORD err = GetLastError();
                log_error( "failed to set end of file: " << err );
                CloseHandle( file );
                return false;
            }

            mem->map = CreateFileMappingA( file, nullptr, PAGE_READWRITE, 0, 0, nullptr );
            CloseHandle( file );

            if ( !mem->map )
            {
                DWORD err = GetLastError();
                log_error( "failed to create new file mapping: " << err );
                return false;
            }

            mem->data = MapViewOfFile( mem->map, FILE_MAP_ALL_ACCESS, 0, 0, 0 );

            if ( !mem->data )
            {
                DWORD err = GetLastError();
                log_error( "failed to map view of expanded file: " << err );
                return false;
            }

            mem->size = new_size;
            log_debug( "file expanded successfully to 0x" << std::hex << mem->size << std::dec );

            dos = static_cast< IMAGE_DOS_HEADER* >( mem->data );
            nt = reinterpret_cast< IMAGE_NT_HEADERS64* >( ( BYTE* ) ( mem->data ) + dos->e_lfanew );
            sec = reinterpret_cast< IMAGE_SECTION_HEADER* >(
                reinterpret_cast< BYTE* >( &nt->OptionalHeader ) +
                nt->FileHeader.SizeOfOptionalHeader );
            last = &sec[ nt->FileHeader.NumberOfSections - 1 ];
        }

        IMAGE_SECTION_HEADER new_sec;
        memset( &new_sec, 0, sizeof( IMAGE_SECTION_HEADER ) );

        memcpy( new_sec.Name, ".edata", 6 );

        new_sec.Misc.VirtualSize = need_size;
        new_sec.VirtualAddress = new_rva;
        new_sec.SizeOfRawData = align_size;
        new_sec.PointerToRawData = new_off;
        new_sec.PointerToRelocations = 0;
        new_sec.PointerToLinenumbers = 0;
        new_sec.NumberOfRelocations = 0;
        new_sec.NumberOfLinenumbers = 0;
        new_sec.Characteristics = IMAGE_SCN_CNT_INITIALIZED_DATA | IMAGE_SCN_MEM_READ;

        memcpy( &sec[ nt->FileHeader.NumberOfSections ], &new_sec, sizeof( IMAGE_SECTION_HEADER ) );
        nt->FileHeader.NumberOfSections++;

        log_debug( "added new section '.edata' at index " << ( nt->FileHeader.NumberOfSections - 1 ) );

        DWORD new_img_size = new_rva + align_size;
        if ( new_img_size > nt->OptionalHeader.SizeOfImage )
        {
            nt->OptionalHeader.SizeOfImage = new_img_size;
            log_debug( "updated SizeOfImage to 0x" << std::hex << new_img_size << std::dec );
        }

        memset( ( BYTE* ) ( mem->data ) + new_off, 0, align_size );

        DWORD exp_off = new_off;
        DWORD exp_rva = new_rva;

        if ( exp_rva % 16 != 0 )
        {
            DWORD a = 16 - ( exp_rva % 16 );
            exp_rva += a;
            exp_off += a;
        }

        log_debug( "placing export directory in new section at RVA 0x" << std::hex << exp_rva
                   << " (offset 0x" << exp_off << ")" << std::dec );

        void* exp_va = ( BYTE* ) ( mem->data ) + exp_off;
        log_debug( "new export directory VA: " << std::hex << exp_va << std::dec );

        auto* new_dir = ( IMAGE_EXPORT_DIRECTORY* ) ( exp_va );

        new_dir->Characteristics = 0;
        new_dir->TimeDateStamp = static_cast< DWORD >( time( nullptr ) );
        new_dir->MajorVersion = 0;
        new_dir->MinorVersion = 0;
        new_dir->Base = 1;
        new_dir->NumberOfFunctions = all_exp.size();
        new_dir->NumberOfNames = all_exp.size();

        DWORD cur_off = exp_off + sizeof( IMAGE_EXPORT_DIRECTORY );
        DWORD cur_rva = exp_rva + sizeof( IMAGE_EXPORT_DIRECTORY );

        std::string dname = path.substr( path.find_last_of( "/\\" ) + 1 );
        log_debug( "setting up DLL name at offset 0x" << std::hex << cur_off << std::dec );
        char* nptr = ( char* ) ( ( BYTE* ) ( mem->data ) + cur_off );
        strcpy_s( nptr, dname.length() + 1, dname.c_str() );
        new_dir->Name = cur_rva;
        cur_off += dname.length() + 1;
        cur_rva += dname.length() + 1;

        if ( cur_off % 4 != 0 )
        {
            DWORD a = 4 - ( cur_off % 4 );
            cur_off += a;
            cur_rva += a;
        }

        log_debug( "setting up functions table at offset 0x" << std::hex << cur_off << std::dec );
        new_dir->AddressOfFunctions = cur_rva;
        DWORD* ftbl = ( DWORD* ) ( ( BYTE* ) ( mem->data ) + cur_off );
        cur_off += all_exp.size() * sizeof( DWORD );
        cur_rva += all_exp.size() * sizeof( DWORD );

        log_debug( "setting up names table at offset 0x" << std::hex << cur_off << std::dec );
        new_dir->AddressOfNames = cur_rva;
        DWORD* ntbl = ( DWORD* ) ( ( BYTE* ) ( mem->data ) + cur_off );
        cur_off += all_exp.size() * sizeof( DWORD );
        cur_rva += all_exp.size() * sizeof( DWORD );

        log_debug( "setting up ordinals table at offset 0x" << std::hex << cur_off << std::dec );
        new_dir->AddressOfNameOrdinals = cur_rva;
        WORD* otbl = ( WORD* ) ( ( BYTE* ) ( mem->data ) + cur_off );
        cur_off += all_exp.size() * sizeof( WORD );
        cur_rva += all_exp.size() * sizeof( WORD );

        for ( size_t i = 0; i < all_exp.size(); i++ )
        {
            log_debug( "processing export " << i << ": " << all_exp[ i ].name
                       << " -> 0x" << std::hex << all_exp[ i ].rva << std::dec );

            ftbl[ i ] = all_exp[ i ].rva;
            otbl[ i ] = static_cast< WORD >( i );

            log_debug( "setting up name string at offset 0x" << std::hex << cur_off << std::dec );
            char* nstr = ( char* ) ( ( BYTE* ) ( mem->data ) + cur_off );
            strcpy_s( nstr, all_exp[ i ].name.length() + 1, all_exp[ i ].name.c_str() );
            ntbl[ i ] = cur_rva;
            cur_off += all_exp[ i ].name.length() + 1;
            cur_rva += all_exp[ i ].name.length() + 1;
        }

        DWORD dir_size = cur_rva - exp_rva;
        log_debug( "new export directory size: 0x" << std::hex << dir_size << std::dec );

        log_debug( "updating data directory" );
        nt->OptionalHeader.DataDirectory[ IMAGE_DIRECTORY_ENTRY_EXPORT ].VirtualAddress = exp_rva;
        nt->OptionalHeader.DataDirectory[ IMAGE_DIRECTORY_ENTRY_EXPORT ].Size = dir_size;

        log_debug( "export directory created in new section successfully" );
        return true;
    }

    bool create_exp_dir()
    {
        log_debug( "creating new export directory" );

        if ( !mem->data )
        {
            log_error( "file data is null" );
            return false;
        }

        std::vector<exp_entry> all = old_exp;
        all.insert( all.end(), new_exp.begin(), new_exp.end() );

        DWORD n = all.size();
        log_debug( "total exports will be: " << n );

        DWORD need = sizeof( IMAGE_EXPORT_DIRECTORY );
        need += sizeof( DWORD ) * n;
        need += sizeof( DWORD ) * n;
        need += sizeof( WORD ) * n;

        std::string dname = path.substr( path.find_last_of( "/\\" ) + 1 );
        need += dname.length() + 1;
        log_debug( "size after adding DLL name (" << dname << "): " << need << " bytes" );

        for ( const auto& e : all )
        {
            need += e.name.length() + 1;
            log_debug( "adding size for export name '" << e.name << "': " << e.name.length() + 1 << " bytes" );
        }

        need = ( need + 15 ) & ~15;
        log_debug( "total needed size (aligned): " << need << " bytes" );

        int rdata_idx = find_rdata_sec();
        if ( rdata_idx == -1 )
        {
            log_error( "could not find .rdata section or suitable section for exports" );
            return false;
        }

        log_debug( "using section index " << rdata_idx << " for exports" );

        IMAGE_SECTION_HEADER& rsec = sec[ rdata_idx ];

        char sname[ 9 ] = { 0 };
        memcpy( sname, rsec.Name, 8 );
        log_debug( "export section name: " << sname );
        log_debug( "export section VirtualAddress: 0x" << std::hex << rsec.VirtualAddress << std::dec );
        log_debug( "export section VirtualSize: 0x" << std::hex << rsec.Misc.VirtualSize << std::dec );
        log_debug( "export section PointerToRawData: 0x" << std::hex << rsec.PointerToRawData << std::dec );
        log_debug( "export section SizeOfRawData: 0x" << std::hex << rsec.SizeOfRawData << std::dec );

        DWORD avail = rsec.SizeOfRawData - rsec.Misc.VirtualSize;

        if ( avail >= need )
        {
            DWORD exp_off = rsec.PointerToRawData + rsec.Misc.VirtualSize;
            DWORD exp_rva = rsec.VirtualAddress + rsec.Misc.VirtualSize;

            if ( exp_rva % 16 != 0 )
            {
                DWORD a = 16 - ( exp_rva % 16 );
                exp_rva += a;
                exp_off += a;
            }

            log_debug( "placing export directory at RVA 0x" << std::hex << exp_rva
                       << " (offset 0x" << exp_off << ")" << std::dec );

            if ( exp_off + need > rsec.PointerToRawData + rsec.SizeOfRawData )
            {
                log_error( "export directory would extend beyond section raw data" );
                return create_new_sec( need, all );
            }

            DWORD v_size = exp_rva + need - rsec.VirtualAddress;
            log_debug( "new virtual size for section: 0x" << std::hex << v_size << std::dec );

            if ( v_size > rsec.SizeOfRawData )
            {
                log_error( "new virtual size would exceed raw data size" );
                return create_new_sec( need, all );
            }

            void* exp_va = ( BYTE* ) ( mem->data ) + exp_off;
            log_debug( "new export directory VA: " << std::hex << exp_va << std::dec );

            memset( exp_va, 0, need );

            auto* new_dir = ( IMAGE_EXPORT_DIRECTORY* ) ( exp_va );

            new_dir->Characteristics = 0;
            new_dir->TimeDateStamp = static_cast< DWORD >( time( nullptr ) );
            new_dir->MajorVersion = 0;
            new_dir->MinorVersion = 0;
            new_dir->Base = 1;
            new_dir->NumberOfFunctions = n;
            new_dir->NumberOfNames = n;

            DWORD cur_off = exp_off + sizeof( IMAGE_EXPORT_DIRECTORY );
            DWORD cur_rva = exp_rva + sizeof( IMAGE_EXPORT_DIRECTORY );

            log_debug( "setting up DLL name at offset 0x" << std::hex << cur_off << std::dec );
            char* nptr = ( char* ) ( ( BYTE* ) ( mem->data ) + cur_off );
            strcpy_s( nptr, dname.length() + 1, dname.c_str() );
            new_dir->Name = cur_rva;
            cur_off += dname.length() + 1;
            cur_rva += dname.length() + 1;

            if ( cur_off % 4 != 0 )
            {
                DWORD a = 4 - ( cur_off % 4 );
                cur_off += a;
                cur_rva += a;
            }

            log_debug( "setting up functions table at offset 0x" << std::hex << cur_off << std::dec );
            new_dir->AddressOfFunctions = cur_rva;
            DWORD* ftbl = ( DWORD* ) ( ( BYTE* ) ( mem->data ) + cur_off );
            cur_off += n * sizeof( DWORD );
            cur_rva += n * sizeof( DWORD );

            log_debug( "setting up names table at offset 0x" << std::hex << cur_off << std::dec );
            new_dir->AddressOfNames = cur_rva;
            DWORD* ntbl = ( DWORD* ) ( ( BYTE* ) ( mem->data ) + cur_off );
            cur_off += n * sizeof( DWORD );
            cur_rva += n * sizeof( DWORD );

            log_debug( "setting up ordinals table at offset 0x" << std::hex << cur_off << std::dec );
            new_dir->AddressOfNameOrdinals = cur_rva;
            WORD* otbl = ( WORD* ) ( ( BYTE* ) ( mem->data ) + cur_off );
            cur_off += n * sizeof( WORD );
            cur_rva += n * sizeof( WORD );

            for ( size_t i = 0; i < n; i++ )
            {
                log_debug( "processing export " << i << ": " << all[ i ].name
                           << " -> 0x" << std::hex << all[ i ].rva << std::dec );

                ftbl[ i ] = all[ i ].rva;
                otbl[ i ] = static_cast< WORD >( i );

                log_debug( "setting up name string at offset 0x" << std::hex << cur_off << std::dec );
                char* nstr = ( char* ) ( ( BYTE* ) ( mem->data ) + cur_off );
                strcpy_s( nstr, all[ i ].name.length() + 1, all[ i ].name.c_str() );
                ntbl[ i ] = cur_rva;
                cur_off += all[ i ].name.length() + 1;
                cur_rva += all[ i ].name.length() + 1;
            }

            DWORD dir_size = cur_rva - exp_rva;
            log_debug( "new export directory size: 0x" << std::hex << dir_size << std::dec );

            log_debug( "updating data directory" );
            nt->OptionalHeader.DataDirectory[ IMAGE_DIRECTORY_ENTRY_EXPORT ].VirtualAddress = exp_rva;
            nt->OptionalHeader.DataDirectory[ IMAGE_DIRECTORY_ENTRY_EXPORT ].Size = dir_size;

            if ( rsec.Misc.VirtualSize < v_size )
            {
                log_debug( "increasing section virtual size from 0x" << std::hex
                           << rsec.Misc.VirtualSize << " to 0x" << v_size << std::dec );
                rsec.Misc.VirtualSize = v_size;
            }

            log_debug( "export directory created successfully in existing section" );
            return true;
        }

        log_debug( "not enough space in existing section. Creating new section." );
        return create_new_sec( need, all );
    }

    int find_rdata_sec() const
    {
        log_debug( "finding .rdata section or suitable section for exports" );

        for ( int i = 0; i < nt->FileHeader.NumberOfSections; i++ )
        {
            char name[ 9 ] = { 0 };
            memcpy( name, sec[ i ].Name, 8 );
            log_debug( "checking section " << i << ": " << name );

            if ( _stricmp( name, ".rdata" ) == 0 )
            {
                log_debug( "found .rdata section at index " << i );
                return i;
            }
        }

        for ( int i = 0; i < nt->FileHeader.NumberOfSections; i++ )
        {
            if ( ( sec[ i ].Characteristics & IMAGE_SCN_CNT_INITIALIZED_DATA ) &&
                 ( sec[ i ].Characteristics & IMAGE_SCN_MEM_READ ) &&
                 !( sec[ i ].Characteristics & IMAGE_SCN_MEM_WRITE ) )
            {
                log_debug( "found suitable read-only data section at index " << i );
                return i;
            }
        }

        log_debug( "no suitable section found" );
        return -1;
    }

    void print_exports() const
    {
        log_debug( "printing exports" );

        if ( !mem->data )
        {
            log_debug( "file data is null, cannot show exports" );
            return;
        }

        DWORD exp_rva = nt->OptionalHeader.DataDirectory[ IMAGE_DIRECTORY_ENTRY_EXPORT ].VirtualAddress;
        log_debug( "current export directory RVA: 0x" << std::hex << exp_rva << std::dec );

        if ( exp_rva == 0 )
        {
            log_info( "no export directory found." );
            return;
        }

        DWORD exp_off = rva_to_off( exp_rva );
        if ( exp_off == 0 && exp_rva != 0 )
        {
            log_error( "could not convert export directory RVA to file offset" );
            return;
        }

        auto* dir = ( IMAGE_EXPORT_DIRECTORY* ) ( ( BYTE* ) ( mem->data ) + exp_off );
        log_debug( "export directory VA: " << std::hex << dir << std::dec );

        if ( ( BYTE* ) ( dir ) < ( BYTE* ) ( mem->data ) ||
             ( BYTE* ) ( dir ) >= ( BYTE* ) ( mem->data ) + mem->size ||
             exp_off + sizeof( IMAGE_EXPORT_DIRECTORY ) > mem->size )
        {
            log_error( "export directory is out of bounds" );
            return;
        }

        DWORD f_off = rva_to_off( dir->AddressOfFunctions );
        DWORD n_off = rva_to_off( dir->AddressOfNames );
        DWORD o_off = rva_to_off( dir->AddressOfNameOrdinals );

        if ( f_off == 0 || n_off == 0 || o_off == 0 )
        {
            log_error( "failed to get offsets for export tables" );
            return;
        }

        DWORD* funcs = ( DWORD* ) ( ( BYTE* ) ( mem->data ) + f_off );
        DWORD* names = ( DWORD* ) ( ( BYTE* ) ( mem->data ) + n_off );
        WORD* ords = ( WORD* ) ( ( BYTE* ) ( mem->data ) + o_off );

        log_debug( "functions table VA: " << std::hex << funcs << std::dec );
        log_debug( "names table VA: " << std::hex << names << std::dec );
        log_debug( "ordinals table VA: " << std::hex << ords << std::dec );

        if ( ( BYTE* ) ( funcs ) < ( BYTE* ) ( mem->data ) ||
             ( BYTE* ) ( funcs ) >= ( BYTE* ) ( mem->data ) + mem->size ||
             ( BYTE* ) ( names ) < ( BYTE* ) ( mem->data ) ||
             ( BYTE* ) ( names ) >= ( BYTE* ) ( mem->data ) + mem->size ||
             ( BYTE* ) ( ords ) < ( BYTE* ) ( mem->data ) ||
             ( BYTE* ) ( ords ) >= ( BYTE* ) ( mem->data ) + mem->size )
        {
            log_error( "export tables are out of bounds" );
            return;
        }

        DWORD name_off = rva_to_off( dir->Name );
        char* dname = ( char* ) ( ( BYTE* ) ( mem->data ) + name_off );

        std::cout << "Export Directory Information:\n";
        std::cout << "  Characteristics: 0x" << std::hex << dir->Characteristics << '\n';
        std::cout << "  TimeDateStamp: 0x" << dir->TimeDateStamp << '\n';
        std::cout << "  MajorVersion: " << std::dec << dir->MajorVersion << '\n';
        std::cout << "  MinorVersion: " << dir->MinorVersion << '\n';

        if ( dname && ( BYTE* ) ( dname ) >= ( BYTE* ) ( mem->data ) &&
             ( BYTE* ) ( dname ) < ( BYTE* ) ( mem->data ) + mem->size )
        {
            std::cout << "  Name: " << dname << '\n';
        }
        else
        {
            std::cout << "  Name: <invalid>\n";
        }

        std::cout << "  Base: " << dir->Base << '\n';
        std::cout << "  NumberOfFunctions: " << dir->NumberOfFunctions << '\n';
        std::cout << "  NumberOfNames: " << dir->NumberOfNames << '\n';
        std::cout << "  Exported Functions:\n";

        for ( DWORD i = 0; i < dir->NumberOfNames; i++ )
        {
            if ( i >= dir->NumberOfFunctions )
            {
                log_debug( "export index " << i << " exceeds function count" );
                break;
            }

            DWORD name_rva = names[ i ];
            WORD ord = ords[ i ];

            if ( ord >= dir->NumberOfFunctions )
            {
                log_debug( "export " << i << ": ordinal " << ord << " is out of range" );
                continue;
            }

            DWORD func_rva = funcs[ ord ];

            DWORD name_off = rva_to_off( name_rva );
            if ( name_off == 0 && name_rva != 0 )
            {
                log_debug( "export " << i << ": failed to get name offset for RVA 0x"
                           << std::hex << name_rva << std::dec );
                continue;
            }

            char* exp_name = ( char* ) ( ( BYTE* ) ( mem->data ) + name_off );

            if ( ( BYTE* ) ( exp_name ) < ( BYTE* ) ( mem->data ) ||
                 ( BYTE* ) ( exp_name ) >= ( BYTE* ) ( mem->data ) + mem->size )
            {
                log_debug( "export " << i << ": name pointer is out of bounds" );
                continue;
            }

            std::cout << "    " << exp_name << " -> "
                << "0x" << std::hex << func_rva << std::dec << '\n';
        }

        log_debug( "finished printing exports" );
    }

    DWORD64 get_image_base() const
    {
        if ( !nt )
        {
            log_error( "nt headers not set" );
            return 0;
        }
        return nt->OptionalHeader.ImageBase;
    }
};

std::string gen_rand_name( int min_len, int max_len )
{
    static const char chars[] = "0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz";

    std::random_device rd;
    std::mt19937 gen( rd() );
    std::uniform_int_distribution len_dist( min_len, max_len );
    std::uniform_int_distribution<> char_dist( 0, sizeof( chars ) - 2 );

    int len = len_dist( gen );
    std::string result;
    result.reserve( len );

    for ( int i = 0; i < len; ++i )
        result += chars[ char_dist( gen ) ];

    return result;
}

DWORD64 gen_rand_va( DWORD64 min_va, DWORD64 max_va )
{
    std::random_device rd;
    std::mt19937_64 gen( rd() );
    std::uniform_int_distribution dist( min_va, max_va );
    return dist( gen );
}

struct args_parser
{
    std::unordered_map<std::string_view, std::string_view> option_map;
    std::vector<std::string_view> positional_args;

    args_parser( int argc, char* argv[] )
    {
        for ( int it = 1; it < argc; it++ )
        {
            std::string_view arg = argv[ it ];

            if ( arg.starts_with( "--" ) )
            {
                auto pos = arg.find( '=' );
                if ( pos != std::string_view::npos )
                {
                    std::string_view key = arg.substr( 0, pos );
                    std::string_view value = arg.substr( pos + 1 );
                    option_map[ key ] = value;
                }
                else if ( it + 1 < argc && !std::string_view( argv[ it + 1 ] ).starts_with( "-" ) )
                {
                    option_map[ arg ] = argv[ ++it ];
                }
                else
                {
                    option_map[ arg ] = "";
                }
            }
            else if ( arg.starts_with( "-" ) && arg.size() > 1 )
            {
                if ( it + 1 < argc && !std::string_view( argv[ it + 1 ] ).starts_with( "-" ) )
                {
                    option_map[ arg ] = argv[ ++it ];
                }
                else
                {
                    option_map[ arg ] = "";
                }
            }
            else
            {
                positional_args.push_back( arg );
            }
        }
    }

    bool has_flag( std::string_view flag ) const
    {
        return option_map.contains( flag );
    }

    std::optional<std::string_view> get_option( std::string_view name ) const
    {
        auto it = option_map.find( name );
        if ( it != option_map.end() )
        {
            return it->second;
        }
        return std::nullopt;
    }

    std::optional<std::string_view> get_arg( size_t index ) const
    {
        if ( index < positional_args.size() )
        {
            return positional_args[ index ];
        }
        return std::nullopt;
    }

    std::optional<uint64_t> get_address( std::string_view name ) const
    {
        auto value = get_option( name );
        if ( !value || value->empty() )
        {
            return std::nullopt;
        }

        uint64_t result = 0;
        std::string_view str = *value;

        if ( str.size() > 2 && str.substr( 0, 2 ) == "0x" )
        {
            str.remove_prefix( 2 );
            auto [ptr, ec] = std::from_chars( str.data(), str.data() + str.size(), result, 16 );
            if ( ec != std::errc() )
            {
                return std::nullopt;
            }
        }
        else
        {
            auto [ptr, ec] = std::from_chars( str.data(), str.data() + str.size(), result, 10 );
            if ( ec != std::errc() )
            {
                return std::nullopt;
            }
        }

        return result;
    }

    std::optional<int> get_int( std::string_view name ) const
    {
        const auto value = get_option( name );
        if ( !value || value->empty() )
        {
            return std::nullopt;
        }

        int result = 0;
        auto [ptr, ec] = std::from_chars( value->data(), value->data() + value->size(), result, 10 );
        if ( ec != std::errc() )
        {
            return std::nullopt;
        }

        return result;
    }

    static uint64_t parse_address( std::string_view str )
    {
        uint64_t result = 0;

        if ( str.size() > 2 && str.substr( 0, 2 ) == "0x" )
        {
            str.remove_prefix( 2 );
            auto [ptr, ec] = std::from_chars( str.data(), str.data() + str.size(), result, 16 );
            if ( ec != std::errc() )
            {
                return 0;
            }
        }
        else
        {
            auto [ptr, ec] = std::from_chars( str.data(), str.data() + str.size(), result, 10 );
            if ( ec != std::errc() )
            {
                return 0;
            }
        }

        return result;
    }
};

void print_usage()
{
    std::cout << R"(
    EXPMOD - just adds exports to 64-bit PE's
    Usage:
      exp <file_path> <export_name> <export_address> [--va | --rva]
      exp <file_path> --random [count]

    Options:
      <file_path>           Path to the 64-bit PE file to modify
      <export_name>         Name of the export to add
      <export_address>      Address of the export (hex with 0x prefix or decimal)
      --va                  Specify that export_address is a Virtual Address (default)
      --rva                 Specify that export_address is a Relative Virtual Address
      --random, -r          Add random exports
      [count]               Number of random exports to add (optional, default: 512)
      --min-va=<addr>       Minimum virtual address for random exports (hex or decimal)
      --max-va=<addr>       Maximum virtual address for random exports (hex or decimal)
      --min-rva=<addr>      Minimum RVA for random exports (hex or decimal)
      --max-rva=<addr>      Maximum RVA for random exports (hex or decimal)

    Example:
      exp mydll.dll MyExport 0x140001000
      exp mydll.dll MyExport 0x1000 --rva
      exp mydll.dll --random 1000 --min-va=0x140010000 --max-va=0x140020000
      exp mydll.dll --random 100 --min-rva=0x1000 --max-rva=0x2000
    )";
}

int main( int argc, char* argv[] )
{
    log_debug( "Starting pe_export_adder" );

    args_parser args( argc, argv );

    auto file_path = args.get_arg( 0 );
    if ( !file_path )
    {
        print_usage();
        return 1;
    }

    expmod e;
    if ( !e.load_file( std::string( *file_path ) ) )
    {
        log_error( "Failed to load file" );
        return 1;
    }

    bool ok = false;
    DWORD64 img_base = e.get_image_base();

    if ( args.has_flag( "--random" ) || args.has_flag( "-r" ) )
    {
        // this is just for examples sake... change it or remove it
        //
        int count = 512;
        DWORD64 min_va = 0x140001000;
        DWORD64 max_va = 0x140003000;

        auto count_opt = args.get_option( "--random" );
        if ( !count_opt ) count_opt = args.get_option( "-r" );

        if ( count_opt && !count_opt->empty() && !count_opt->starts_with( "-" ) )
        {
            try
            {
                count = std::stoi( std::string( *count_opt ) );
                if ( count <= 0 ) count = 512;
            }
            catch ( ... )
            {
            }
        }
        else if ( auto count_arg = args.get_arg( 1 ) )
        {
            if ( !count_arg->starts_with( "-" ) )
            {
                try
                {
                    count = std::stoi( std::string( *count_arg ) );
                    if ( count <= 0 ) count = 512;
                }
                catch ( ... )
                {
                }
            }
        }

        if ( auto val = args.get_address( "--min-va" ) )
        {
            min_va = *val;
        }
        if ( auto val = args.get_address( "--max-va" ) )
        {
            max_va = *val;
        }
        if ( auto val = args.get_address( "--min-rva" ) )
        {
            min_va = *val + img_base;
        }
        if ( auto val = args.get_address( "--max-rva" ) )
        {
            max_va = *val + img_base;
        }

        log_info( "adding " << count << " random exports with VA range 0x"
                  << std::hex << min_va << " - 0x" << max_va << std::dec );

        for ( int i = 0; i < count; i++ )
        {
            std::string name = gen_rand_name( 8, 16 );
            DWORD64 addr = gen_rand_va( min_va, max_va );

            if ( !e.add_export( name, addr ) )
            {
                log_error( "failed to add random export #" << ( i + 1 ) );
                return 1;
            }
        }

        ok = true;
    }
    else
    {
        auto name = args.get_arg( 1 );
        auto addr_str = args.get_arg( 2 );

        if ( name && addr_str )
        {
            try
            {
                DWORD64 addr = args.parse_address( *addr_str );
                bool is_rva = args.has_flag( "--rva" );

                if ( is_rva )
                {
                    addr += img_base;
                    log_info( "adding export: " << *name << " -> RVA 0x" << std::hex
                              << ( addr - img_base ) << " (VA 0x" << addr << ")" << std::dec );
                }
                else
                {
                    log_info( "adding export: " << *name << " -> VA 0x" << std::hex
                              << addr << " (RVA 0x" << ( addr - img_base ) << ")" << std::dec );
                }

                if ( !e.add_export( std::string( *name ), addr ) )
                {
                    log_error( "failed to add export" );
                    return 1;
                }

                ok = true;
            }
            catch ( const std::exception& e )
            {
                log_error( "error parsing address: " << e.what() );
                return 1;
            }
        }
        else
        {
            print_usage();
            return 1;
        }
    }

    if ( ok )
    {
        if ( !e.update_exports() )
        {
            log_error( "failed to update export directory" );
            return 1;
        }

        e.print_exports();
        log_info( "exports successfully added to " << *file_path );
    }

    return 0;
}