Trying to create a PE executable using cpp

portable.cpp

#include <cstdint>
#include <cstring>
#include <ctime>
#include <fstream>
#include <iostream>
#include <vector>

#pragma pack(push, 1)

// DOS Header (64 bytes)
struct DOSHeader
{
    uint16_t e_magic;    // Magic number
    uint16_t e_cblp;     // Bytes on last page of file
    uint16_t e_cp;       // Pages in file
    uint16_t e_crlc;     // Relocations
    uint16_t e_cparhdr;  // Size of header in paragraphs
    uint16_t e_minalloc; // Minimum extra paragraphs needed
    uint16_t e_maxalloc; // Maximum extra paragraphs needed
    uint16_t e_ss;       // Initial (relative) SS value
    uint16_t e_sp;       // Initial SP value
    uint16_t e_csum;     // Checksum
    uint16_t e_ip;       // Initial IP value
    uint16_t e_cs;       // Initial (relative) CS value
    uint16_t e_lfarlc;   // File address of relocation table
    uint16_t e_ovno;     // Overlay number
    uint16_t e_res[4];   // Reserved words
    uint16_t e_oemid;    // OEM identifier (for e_oeminfo)
    uint16_t e_oeminfo;  // OEM information; e_oemid specific
    uint16_t e_res2[10]; // Reserved words
    int32_t e_lfanew;    // File address of new exe header
};

// COFF File Header (20 bytes)
struct COFFHeader
{
    uint16_t Machine;
    uint16_t NumberOfSections;
    uint32_t TimeDateStamp;
    uint32_t PointerToSymbolTable;
    uint32_t NumberOfSymbols;
    uint16_t SizeOfOptionalHeader;
    uint16_t Characteristics;
};

// PE Optional Header (240 bytes)
struct OptionalHeader
{
    uint16_t Magic;
    uint8_t MajorLinkerVersion;
    uint8_t MinorLinkerVersion;
    uint32_t SizeOfCode;
    uint32_t SizeOfInitializedData;
    uint32_t SizeOfUninitializedData;
    uint32_t AddressOfEntryPoint;
    uint32_t BaseOfCode;
    uint64_t ImageBase;
    uint32_t SectionAlignment;
    uint32_t FileAlignment;
    uint16_t MajorOperatingSystemVersion;
    uint16_t MinorOperatingSystemVersion;
    uint16_t MajorImageVersion;
    uint16_t MinorImageVersion;
    uint16_t MajorSubsystemVersion;
    uint16_t MinorSubsystemVersion;
    uint32_t Win32VersionValue;
    uint32_t SizeOfImage;
    uint32_t SizeOfHeaders;
    uint32_t CheckSum;
    uint16_t Subsystem;
    uint16_t DllCharacteristics;
    uint64_t SizeOfStackReserve;
    uint64_t SizeOfStackCommit;
    uint64_t SizeOfHeapReserve;
    uint64_t SizeOfHeapCommit;
    uint32_t LoaderFlags;
    uint32_t NumberOfRvaAndSizes;
    uint64_t DataDirectory[16];
};

// Section Header (40 bytes)
struct SectionHeader
{
    uint8_t Name[8];
    uint32_t VirtualSize;
    uint32_t VirtualAddress;
    uint32_t SizeOfRawData;
    uint32_t PointerToRawData;
    uint32_t PointerToRelocations;
    uint32_t PointerToLinenumbers;
    uint16_t NumberOfRelocations;
    uint16_t NumberOfLinenumbers;
    uint32_t Characteristics;
};

#pragma pack(pop)

// Define a structure to represent an instruction
struct Instruction
{
    std::string mnemonic;
    std::vector<std::string> operands;
};

// Function to parse a line of assembly code
Instruction parseInstruction(const std::string &line)
{
    Instruction instruction;
    size_t pos = line.find(' ');
    instruction.mnemonic = line.substr(0, pos);
    if (pos != std::string::npos) {
        std::string operands_str = line.substr(pos + 1);
        size_t start = 0, end = 0;
        while ((end = operands_str.find(',', start)) != std::string::npos) {
            instruction.operands.push_back(operands_str.substr(start, end - start));
            start = end + 1;
        }
        instruction.operands.push_back(operands_str.substr(start));
    }
    return instruction;
}

// Function to assemble the parsed instructions into machine code
std::vector<uint8_t> assemble(const std::vector<Instruction> &instructions)
{
    std::vector<uint8_t> machineCode;

    for (const auto &instruction : instructions) {
        if (instruction.mnemonic == "mov") {
            if (instruction.operands[0] == "rax" && instruction.operands[1] == "1") {
                // mov rax, 1
                machineCode.push_back(0x48);
                machineCode.push_back(0xC7);
                machineCode.push_back(0xC0);
                machineCode.push_back(0x01);
                machineCode.push_back(0x00);
                machineCode.push_back(0x00);
                machineCode.push_back(0x00);
            }
        } else if (instruction.mnemonic == "syscall") {
            // syscall
            machineCode.push_back(0x0F);
            machineCode.push_back(0x05);
        }
    }

    return machineCode;
}

class PEFile
{
public:
    enum class Type { Executable, Library };

    PEFile(const std::string &filename, Type type);
    void addCodeSection(const std::vector<uint8_t> &code);
    void addResourceSection(const std::vector<uint8_t> &resources);
    void writeToFile();

private:
    std::string filename;
    Type type;
    DOSHeader dosHeader;
    COFFHeader coffHeader;
    OptionalHeader optionalHeader;
    SectionHeader textSectionHeader;
    SectionHeader rsrcSectionHeader;
    std::vector<uint8_t> codeSection;
    std::vector<uint8_t> resourceSection;
    size_t align(size_t size, size_t alignment);
    void calculateHeaders();
    void initializeHeaders();
};

PEFile::PEFile(const std::string &filename, Type type)
    : filename(filename)
    , type(type)
{
    initializeHeaders();
}

void PEFile::initializeHeaders()
{
    // Initialize DOS Header
    dosHeader = {};
    dosHeader.e_magic = 0x5A4D; // MZ
    dosHeader.e_cblp = 0x0090;
    dosHeader.e_cp = 0x0003;
    dosHeader.e_crlc = 0x0000;
    dosHeader.e_cparhdr = 0x0004;
    dosHeader.e_minalloc = 0x0000;
    dosHeader.e_maxalloc = 0xFFFF;
    dosHeader.e_ss = 0x0000;
    dosHeader.e_sp = 0x00B8;
    dosHeader.e_csum = 0x0000;
    dosHeader.e_ip = 0x0000;
    dosHeader.e_cs = 0x0000;
    dosHeader.e_lfarlc = 0x0040;
    dosHeader.e_ovno = 0x0000;
    dosHeader.e_res[0] = 0x0000;
    dosHeader.e_res[1] = 0x0000;
    dosHeader.e_res[2] = 0x0000;
    dosHeader.e_res[3] = 0x0000;
    dosHeader.e_oemid = 0x0000;
    dosHeader.e_oeminfo = 0x0000;
    dosHeader.e_res2[0] = 0x0000;
    dosHeader.e_res2[1] = 0x0000;
    dosHeader.e_res2[2] = 0x0000;
    dosHeader.e_res2[3] = 0x0000;
    dosHeader.e_res2[4] = 0x0000;
    dosHeader.e_res2[5] = 0x0000;
    dosHeader.e_res2[6] = 0x0000;
    dosHeader.e_res2[7] = 0x0000;
    dosHeader.e_res2[8] = 0x0000;
    dosHeader.e_res2[9] = 0x0000;
    dosHeader.e_lfanew = sizeof(DOSHeader);

    // Initialize COFF Header
    coffHeader = {};
    coffHeader.Machine = 0x8664; // x86-64
    coffHeader.NumberOfSections = 2;
    coffHeader.TimeDateStamp = time(nullptr); // Current timestamp
    coffHeader.PointerToSymbolTable = 0x00000000;
    coffHeader.NumberOfSymbols = 0;
    coffHeader.SizeOfOptionalHeader = sizeof(OptionalHeader);

    // Set Characteristics and Optional Header based on the type
    if (type == Type::Executable) {
        coffHeader.Characteristics = 0x0002 | 0x0100
                                     | 0x0020; // Executable Image, 32-bit machine, No line numbers
        optionalHeader.Subsystem = 3;          // CUI (Console)
        optionalHeader.DllCharacteristics = 0;
    } else if (type == Type::Library) {
        coffHeader.Characteristics = 0x2000 | 0x0004
                                     | 0x0020;      // DLL, 32-bit machine, No line numbers
        optionalHeader.Subsystem = 2;               // Windows GUI
        optionalHeader.DllCharacteristics = 0x0040; // Dynamic base
    }

    // Initialize Optional Header
    optionalHeader = {};
    optionalHeader.Magic = 0x20B; // PE32+ or use 0x10B for 32bit
    optionalHeader.MajorLinkerVersion = 1;
    optionalHeader.MinorLinkerVersion = 0;
    optionalHeader.AddressOfEntryPoint = 0x00001000;
    optionalHeader.BaseOfCode = 0x00001000;
    optionalHeader.ImageBase = 0x000400000;
    optionalHeader.SectionAlignment = 0x00001000;
    optionalHeader.FileAlignment = 0x00000200;
    optionalHeader.MajorOperatingSystemVersion = 6;
    optionalHeader.MinorOperatingSystemVersion = 0;
    optionalHeader.MajorImageVersion = 0;
    optionalHeader.MinorImageVersion = 1;
    optionalHeader.MajorSubsystemVersion = 6;
    optionalHeader.MinorSubsystemVersion = 0;
    optionalHeader.Win32VersionValue = 0;
    optionalHeader.SizeOfImage = 0; // To be calculated
    optionalHeader.SizeOfHeaders = 0x00000200;
    optionalHeader.CheckSum = 0; // Should be calculated
    optionalHeader.SizeOfStackReserve = 0x0000000010000000;
    optionalHeader.SizeOfStackCommit = 0x0000000000001000;
    optionalHeader.SizeOfHeapReserve = 0x0000000010000000;
    optionalHeader.SizeOfHeapCommit = 0x0000000000001000;
    optionalHeader.LoaderFlags = 0;
    optionalHeader.NumberOfRvaAndSizes = 16;
}

void PEFile::addCodeSection(const std::vector<uint8_t> &code)
{
    codeSection = code;
    textSectionHeader = {};
    std::memcpy(textSectionHeader.Name, ".text", 5);
    textSectionHeader.VirtualSize = code.size();
    textSectionHeader.VirtualAddress = 0x00001000;
    textSectionHeader.SizeOfRawData = align(code.size(), optionalHeader.FileAlignment);
    textSectionHeader.PointerToRawData = 0x00000200;
    textSectionHeader.PointerToRelocations = 0x00000000;
    textSectionHeader.PointerToLinenumbers = 0x00000000;
    textSectionHeader.NumberOfRelocations = 0;
    textSectionHeader.NumberOfLinenumbers = 0;
    textSectionHeader.Characteristics = 0x60000020; // Code, Execute, Read
}

void PEFile::addResourceSection(const std::vector<uint8_t> &resources)
{
    resourceSection = resources;
    rsrcSectionHeader = {};
    std::memcpy(rsrcSectionHeader.Name, ".rsrc", 5);
    rsrcSectionHeader.VirtualSize = resources.size();
    rsrcSectionHeader.VirtualAddress = 0x00002000;
    rsrcSectionHeader.SizeOfRawData = align(resources.size(), optionalHeader.FileAlignment);
    rsrcSectionHeader.PointerToRawData = textSectionHeader.PointerToRawData
                                         + textSectionHeader.SizeOfRawData;
    rsrcSectionHeader.PointerToRelocations = 0x00000000;
    rsrcSectionHeader.PointerToLinenumbers = 0x00000000;
    rsrcSectionHeader.NumberOfRelocations = 0;
    rsrcSectionHeader.NumberOfLinenumbers = 0;
    rsrcSectionHeader.Characteristics = 0x40000040; // Initialized data, Read
}

void PEFile::writeToFile()
{
    calculateHeaders();

    std::ofstream file(filename, std::ios::binary);

    // Write DOS header
    file.write(reinterpret_cast<const char *>(&dosHeader), sizeof(dosHeader));

    // Write PE signature
    file.write("PE\0\0", 4);

    // Write COFF header
    file.write(reinterpret_cast<const char *>(&coffHeader), sizeof(coffHeader));

    // Write Optional header
    file.write(reinterpret_cast<const char *>(&optionalHeader), sizeof(optionalHeader));

    // Write section headers
    file.write(reinterpret_cast<const char *>(&textSectionHeader), sizeof(textSectionHeader));
    file.write(reinterpret_cast<const char *>(&rsrcSectionHeader), sizeof(rsrcSectionHeader));

    // Write code section
    file.write(reinterpret_cast<const char *>(codeSection.data()), codeSection.size());

    // Align to section boundary
    std::vector<uint8_t> padding(textSectionHeader.SizeOfRawData - codeSection.size());
    file.write(reinterpret_cast<const char *>(padding.data()), padding.size());

    // Write resource section
    file.write(reinterpret_cast<const char *>(resourceSection.data()), resourceSection.size());

    // Align to section boundary
    padding.resize(rsrcSectionHeader.SizeOfRawData - resourceSection.size());
    file.write(reinterpret_cast<const char *>(padding.data()), padding.size());

    file.close();
    std::cout << "Executable file '" << filename << "' created successfully." << std::endl;
}

size_t PEFile::align(size_t size, size_t alignment)
{
    return (size + alignment - 1) & ~(alignment - 1);
}

void PEFile::calculateHeaders()
{
    optionalHeader.SizeOfCode = textSectionHeader.SizeOfRawData;
    optionalHeader.SizeOfInitializedData = rsrcSectionHeader.SizeOfRawData;
    optionalHeader.SizeOfImage = rsrcSectionHeader.VirtualAddress
                                 + align(rsrcSectionHeader.VirtualSize,
                                         optionalHeader.SectionAlignment);
}

int main()
{
    // Sample assembly code
    std::vector<std::string> assemblyCode = {"mov rax, 1", "syscall"};

    // Parse the assembly code
    std::vector<Instruction> instructions;
    for (const auto &line : assemblyCode) {
        instructions.push_back(parseInstruction(line));
    }

    // Assemble the instructions into machine code
    std::vector<uint8_t> machineCode = assemble(instructions);

    // Create PE file
    PEFile peFile("hello_world.exe", PEFile::Type::Executable);

    // Add sections
    peFile.addCodeSection(machineCode);
    peFile.addResourceSection(std::vector<uint8_t>(0x400, 0)); // Dummy resource data

    // Write to file
    peFile.writeToFile();

    return 0;
}