pe.cpp

#include <cstdlib>
#include <iostream>
#include <zasm/builder/builder.hpp>
#include <zasm/zasm.hpp>

static void generateTestProgram(zasm::Program& program)
{
    using namespace zasm;
    using namespace zasm::operands;

    Assembler a(program);

    Label labelMain = a.createLabel("main");

    Label labelData01 = a.createLabel("data01");
    Label labelData02 = a.createLabel("data02");
    Label labelData03 = a.createLabel("data03");
    Label labelData04 = a.createLabel("data04");
    Label labelData05 = a.createLabel("data05");

    Label labelRData01 = a.createLabel("rdata01");
    Label labelRData02 = a.createLabel("rdata02");
    Label labelRData03 = a.createLabel("rdata03");
    Label labelRData04 = a.createLabel("rdata04");
    Label labelRData05 = a.createLabel("rdata05");

    Label labelStrTitle = a.createLabel("strTitle");
    Label labelStrMessage = a.createLabel("strMessage");

    Label labelImpExitProcess = program.getOrCreateImportLabel("kernel32.dll", "ExitProcess");
    Label labelImpMessageBoxA = program.getOrCreateImportLabel("user32.dll", "MessageBoxA");

    a.section(".text", Section::Attribs::Code | Section::Attribs::Exec | Section::Attribs::Read);
    {
        a.bind(labelMain);
        a.sub(rsp, Imm(0x28));

        // Test read/write on sections.
        {
            a.lea(rdx, qword_ptr(labelRData01));
            a.lea(rax, qword_ptr(rdx));
            a.lea(rdx, qword_ptr(labelData01));
            a.mov(qword_ptr(rdx), rax);

            a.lea(rdx, qword_ptr(labelRData02));
            a.lea(rax, qword_ptr(rdx));
            a.lea(rdx, qword_ptr(labelData02));
            a.mov(qword_ptr(rdx), rax);

            a.lea(rdx, qword_ptr(labelRData03));
            a.lea(rax, qword_ptr(rdx));
            a.lea(rdx, qword_ptr(labelData03));
            a.mov(qword_ptr(rdx), rax);

            a.lea(rdx, qword_ptr(labelRData04));
            a.lea(rax, qword_ptr(rdx));
            a.lea(rdx, qword_ptr(labelData04));
            a.mov(qword_ptr(rdx), rax);

            a.lea(rdx, qword_ptr(labelRData05));
            a.lea(rax, qword_ptr(rdx));
            a.lea(rdx, qword_ptr(labelData05));
            a.mov(qword_ptr(rdx), rax);
        }

        // MessageBox
        {
            a.xor_(rcx, rcx);                       // hWnd
            a.lea(rdx, qword_ptr(labelStrMessage)); // lpText
            a.lea(r8, qword_ptr(labelStrTitle));    // lpCaption
            a.mov(r9, Imm(0));                      // uType
            a.call(qword_ptr(labelImpMessageBoxA));
        }

        // ExitProcess
        {
            a.xor_(rcx, rcx);
            a.call(qword_ptr(labelImpExitProcess));
        }

        a.add(rsp, Imm(0x28));
        a.ret();
    }

    a.section(".data", Section::Attribs::Data | Section::Attribs::Read | Section::Attribs::Write);
    {
        a.bind(labelData01);
        a.dq(01);
        a.bind(labelData02);
        a.dq(02);
        a.bind(labelData03);
        a.dq(03);
        a.bind(labelData04);
        a.dq(04);
        a.bind(labelData05);
        a.dq(05);
    }

    a.section(".rdata", Section::Attribs::RData | Section::Attribs::Read);
    {
        a.bind(labelStrTitle);
        a.embed("Hello", 6);

        a.bind(labelStrMessage);
        a.embed("Hello World", 12);

        a.bind(labelRData01);
        a.dq(01);
        a.bind(labelRData02);
        a.dq(02);
        a.bind(labelRData03);
        a.dq(03);
        a.bind(labelRData04);
        a.dq(04);
        a.bind(labelRData05);
        a.dq(05);
    }

    // Program start.
    program.setEntryPoint(labelMain);
}

int main()
{
    auto program = zasm::Program(ZydisMachineMode::ZYDIS_MACHINE_MODE_LONG_64);

    generateTestProgram(program);

    auto moduleBuilder = zasm::frontend::getModuleBuilder(zasm::frontend::ModuleType::PE);

    if (!moduleBuilder->generate(program, std::filesystem::path("./test.exe")))
    {
        std::cerr << "Failed to generate module\n";
        return EXIT_FAILURE;
    }

    return EXIT_SUCCESS;
}