#include "sandbox.h"
auto Sandbox::BuildPebParameter() -> void {
    //设置peb->ProcessParameters
           // 设置进程参数
    RTL_USER_PROCESS_PARAMETERS processParams = {};

    // 映射进程参数结构内存
    size_t processParamsSize = AlignToSectionAlignment(sizeof(RTL_USER_PROCESS_PARAMETERS), PAGE_SIZE);
    uc_err procErr = uc_mem_map(m_ucEngine, m_processParamsBase, processParamsSize,
        UC_PROT_READ | UC_PROT_WRITE);
    if (procErr != UC_ERR_OK) {
        throw std::runtime_error("Failed to map process parameters block");
    }

    // 设置基本参数
    processParams.MaximumLength = sizeof(RTL_USER_PROCESS_PARAMETERS);
    processParams.Length = sizeof(RTL_USER_PROCESS_PARAMETERS);

    // 字符串缓冲区基址
    uint64_t stringBufferBase = m_processParamsBase + processParamsSize;
    uint64_t currentStringPos = stringBufferBase;

    // 映射字符串缓冲区
    size_t stringBufferSize = AlignToSectionAlignment(4096, PAGE_SIZE); // 足够大的空间存储字符串
    uc_mem_map(m_ucEngine, stringBufferBase, stringBufferSize, UC_PROT_READ | UC_PROT_WRITE);

    // 设置映像路径
    std::wstring imagePath = L"C:\\Path\\To\\EmulatedImage.exe";
    UNICODE_STRING imagePathUnicode;
    imagePathUnicode.Length = static_cast<USHORT>(imagePath.length() * sizeof(wchar_t));
    imagePathUnicode.MaximumLength = imagePathUnicode.Length + sizeof(wchar_t); // 额外空间用于null终止符
    imagePathUnicode.Buffer = reinterpret_cast<PWSTR>(currentStringPos);

    // 写入映像路径字符串
    uc_mem_write(m_ucEngine, currentStringPos, imagePath.c_str(), imagePath.length() * sizeof(wchar_t) + sizeof(wchar_t));
    currentStringPos += AlignToSectionAlignment(imagePathUnicode.MaximumLength, 8);

    // 设置命令行
    std::wstring commandLine = L"EmulatedImage.exe argument1 argument2";
    UNICODE_STRING commandLineUnicode;
    commandLineUnicode.Length = static_cast<USHORT>(commandLine.length() * sizeof(wchar_t));
    commandLineUnicode.MaximumLength = commandLineUnicode.Length + sizeof(wchar_t);
    commandLineUnicode.Buffer = reinterpret_cast<PWSTR>(currentStringPos);

    // 写入命令行字符串
    uc_mem_write(m_ucEngine, currentStringPos, commandLine.c_str(), commandLine.length() * sizeof(wchar_t) + sizeof(wchar_t));
    currentStringPos += AlignToSectionAlignment(commandLineUnicode.MaximumLength, 8);

    // 设置当前目录
    std::wstring currentDir = L"C:\\Path\\To";
    UNICODE_STRING currentDirUnicode;
    currentDirUnicode.Length = static_cast<USHORT>(currentDir.length() * sizeof(wchar_t));
    currentDirUnicode.MaximumLength = currentDirUnicode.Length + sizeof(wchar_t);
    currentDirUnicode.Buffer = reinterpret_cast<PWSTR>(currentStringPos);

    // 写入当前目录字符串
    uc_mem_write(m_ucEngine, currentStringPos, currentDir.c_str(), currentDir.length() * sizeof(wchar_t) + sizeof(wchar_t));
    currentStringPos += AlignToSectionAlignment(currentDirUnicode.MaximumLength, 8);

    // 设置DLL路径
    std::wstring dllPath = L"C:\\Windows\\System32";
    UNICODE_STRING dllPathUnicode;
    dllPathUnicode.Length = static_cast<USHORT>(dllPath.length() * sizeof(wchar_t));
    dllPathUnicode.MaximumLength = dllPathUnicode.Length + sizeof(wchar_t);
    dllPathUnicode.Buffer = reinterpret_cast<PWSTR>(currentStringPos);

    // 写入DLL路径字符串
    uc_mem_write(m_ucEngine, currentStringPos, dllPath.c_str(), dllPath.length() * sizeof(wchar_t) + sizeof(wchar_t));

    // 设置进程参数结构中的字符串字段
    processParams.ImagePathName = imagePathUnicode;
    processParams.CommandLine = commandLineUnicode;
    processParams.CurrentDirectory.DosPath = currentDirUnicode;
    processParams.DllPath = dllPathUnicode;
    processParams.Environment = reinterpret_cast<WCHAR*>(m_envBlockBase);

    // 写入进程参数结构
    uc_mem_write(m_ucEngine, m_processParamsBase, &processParams, sizeof(RTL_USER_PROCESS_PARAMETERS));

}
auto Sandbox::SetupVirtualMachine() -> void {
    SegmentSelector cs = { 0 };
    cs.fields.index = 1;
    uc_reg_write(m_ucEngine, UC_X86_REG_CS, &cs.all);

    SegmentSelector ds = { 0 };
    ds.fields.index = 2;
    uc_reg_write(m_ucEngine, UC_X86_REG_DS, &ds.all);

    SegmentSelector ss = { 0 };
    ss.fields.index = 2;
    uc_reg_write(m_ucEngine, UC_X86_REG_SS, &ss.all);

    SegmentSelector es = { 0 };
    es.fields.index = 2;
    uc_reg_write(m_ucEngine, UC_X86_REG_ES, &es.all);

    SegmentSelector gs = { 0 };
    gs.fields.index = 2;
    uc_reg_write(m_ucEngine, UC_X86_REG_GS, &gs.all);

    FlagRegister eflags = { 0 };
    eflags.fields.id = 1;
    eflags.fields.intf = 1;
    eflags.fields.reserved1 = 1;

    uc_reg_write(m_ucEngine, UC_X86_REG_EFLAGS, &eflags.all);

    uint64_t cr8 = 0;
    uc_reg_write(m_ucEngine, UC_X86_REG_CR8, &cr8);

    /*
        映射 m_KSharedUserDataBase
    */
    m_KSharedUserDataBase = 0x7FFE0000;
    uint64_t m_KSharedUserDataEnd = 0x7FFE0FFF;  // 0x7FFE2000
    m_KSharedUserDataSize = AlignToSectionAlignment(
        m_KSharedUserDataEnd - m_KSharedUserDataBase, PAGE_SIZE);

    uc_mem_map(m_ucEngine, m_KSharedUserDataBase, m_KSharedUserDataSize,
        UC_PROT_READ);
    uc_mem_write(m_ucEngine, m_KSharedUserDataBase,
        (void*)m_KSharedUserDataBase, m_KSharedUserDataSize);

    m_tebBase = TEB_BASE;             // 进程TEB地址
    m_pebBase = PEB_BASE;             // 进程PEB地址
    m_envBlockBase = ENV_BLOCK_BASE;  // 环境变量块地址
    m_processParamsBase = PROCESS_PARAMS_BASE; // 进程参数地址

    // stack
    m_stackBase = AlignToSectionAlignment(
        this->m_peInfo->isX64 ? STACK_BASE_64 : STACK_BASE_32, 16);
    m_stackSize = AlignToSectionAlignment(
        this->m_peInfo->isX64 ? STACK_SIZE_64 : STACK_SIZE_32, 16);
    m_stackEnd = m_stackBase + m_stackSize;

    // heap
    m_heapBase = this->m_peInfo->isX64 ? HEAP_ADDRESS_64 : HEAP_ADDRESS_32;
    m_heapSize = this->m_peInfo->isX64 ? HEAP_SIZE_64 : HEAP_SIZE_32;
    m_heapEnd = m_heapBase + m_heapSize;

    BuildPebParameter();
    InitializeLdrData();

    // 根据PE文件类型设置PEB和TEB
    if (this->m_peInfo->isX64) {
        // 设置64位PEB
        m_peb64.ImageBaseAddress = m_peInfo->RecImageBase;
        m_pebEnd =
            m_pebBase + AlignToSectionAlignment(sizeof(X64PEB), PAGE_SIZE);
        m_tebEnd =
            m_tebBase + AlignToSectionAlignment(sizeof(X64TEB), PAGE_SIZE);

        // 设置64位TEB
        m_teb64.ClientId.UniqueProcess = GetCurrentProcessId();
        m_teb64.ClientId.UniqueThread = GetCurrentThreadId();
        m_teb64.ProcessEnvironmentBlock = reinterpret_cast<X64PEB*>(m_pebBase);
        m_teb64.NtTib.StackBase = (DWORD64)m_stackBase;
        m_teb64.NtTib.StackLimit = (DWORD64)m_stackSize;
        m_peb64.ProcessParameters = m_processParamsBase;

        // 设置堆
        m_peb64.ProcessHeap = m_heapBase;

        // 设置GS基址结构
        m_gsBaseStruct.teb = m_tebBase;
        m_gsBaseStruct.peb = m_pebBase;
        uint64_t gsAllocSize =
            AlignToSectionAlignment(sizeof(struct_gs_base), PAGE_SIZE);

        // 映射PEB到虚拟内存
        uc_mem_map(m_ucEngine, m_pebBase, m_pebEnd - m_pebBase,
            UC_PROT_READ | UC_PROT_WRITE);
        uc_mem_write(m_ucEngine, m_pebBase, &m_peb64, sizeof(X64PEB));

        // 映射TEB到虚拟内存
        uc_mem_map(m_ucEngine, m_tebBase, m_tebEnd - m_tebBase,
            UC_PROT_READ | UC_PROT_WRITE);
        uc_mem_write(m_ucEngine, m_tebBase, &m_teb64, sizeof(X64TEB));

        // 映射GS基址结构到虚拟内存
        uc_mem_map(m_ucEngine, m_gsBase, gsAllocSize, UC_PROT_READ);
        uc_mem_write(m_ucEngine, m_gsBase, &m_gsBaseStruct,
            sizeof(struct_gs_base));

        // 设置GS基址MSR
        uc_x86_msr msr;
        msr.rid = static_cast<uint32_t>(Msr::kIa32GsBase);
        msr.value = m_gsBase;
        uc_reg_write(m_ucEngine, UC_X86_REG_MSR, &msr);
    }
    else {
        // 设置32位PEB
        m_peb32.ImageBaseAddress = static_cast<ULONG>(m_peInfo->RecImageBase);
        m_pebEnd =
            m_pebBase + AlignToSectionAlignment(sizeof(X32PEB), PAGE_SIZE);
        m_tebEnd =
            m_tebBase + AlignToSectionAlignment(sizeof(X32TEB), PAGE_SIZE);

        // 设置32位TEB
        m_teb32.ClientId.UniqueProcess = GetCurrentProcessId();
        m_teb32.ClientId.UniqueThread = GetCurrentThreadId();
        m_teb32.ProcessEnvironmentBlock = static_cast<ULONG>(m_pebBase);
        m_teb32.NtTib.StackBase = static_cast<ULONG>(m_stackBase);
        m_teb32.NtTib.StackLimit = static_cast<ULONG>(m_stackSize);
        // 初始化NT_TIB结构的其余部分
        m_teb32.NtTib.Self =
            static_cast<ULONG>(m_tebBase);  // 关键：设置Self指针指向TEB本身
        m_teb32.NtTib.ExceptionList = 0xFFFFFFFF;  // 初始异常链表指向特殊值
        m_teb32.NtTib.Version = 0;
        m_teb32.NtTib.FiberData = 0;
        m_teb32.NtTib.ArbitraryUserPointer = 0;
        m_peb32.ProcessParameters = m_processParamsBase;

        // 设置堆
        m_peb32.ProcessHeap = static_cast<ULONG>(m_heapBase);

        // 映射PEB到虚拟内存
        uc_mem_map(m_ucEngine, m_pebBase, m_pebEnd - m_pebBase,
            UC_PROT_READ | UC_PROT_WRITE);
        uc_mem_write(m_ucEngine, m_pebBase, &m_peb32, sizeof(X32PEB));

        // 映射TEB到虚拟内存
        uc_mem_map(m_ucEngine, m_tebBase, m_tebEnd - m_tebBase,
            UC_PROT_READ | UC_PROT_WRITE);
        uc_mem_write(m_ucEngine, m_tebBase, &m_teb32, sizeof(X32TEB));

        // 对于32位，我们需要设置FS段寄存器指向TEB
        SegmentSelector fs = { 0 };
        fs.fields.index = 3;
        // 不需要设置present和dpl，因为SegmentSelector结构体中没有这些字段
        uc_reg_write(m_ucEngine, UC_X86_REG_FS, &fs.all);

        // 设置FS基址MSR
        uc_x86_msr msr;
        msr.rid = static_cast<uint32_t>(Msr::kIa32FsBase);
        msr.value = m_tebBase;
        uc_reg_write(m_ucEngine, UC_X86_REG_MSR, &msr);

        // 确保TEB中关键字段被正确初始化
        // 特别是FS:18h (0x18)处应该指向自身
        // 根据Native_Struct.h中X32TEB定义，偏移0x18处是SelfTeb
        uint32_t self_teb_ptr = static_cast<uint32_t>(m_tebBase);
        // 在NT_TIB中设置SelfTeb (offset 0x18)
        uc_mem_write(m_ucEngine, m_tebBase + 0x18, &self_teb_ptr,
            sizeof(uint32_t));

        // 确保TEB中的ProcessEnvironmentBlock字段指向PEB
        uint32_t peb_ptr = static_cast<uint32_t>(m_pebBase);
        // 偏移0x30处是ProcessEnvironmentBlock
        uc_mem_write(m_ucEngine, m_tebBase + 0x30, &peb_ptr, sizeof(uint32_t));
    }
    // 映射新的内存区域
    size_t envSize =
        AlignToSectionAlignment(this->GetEnvStringsSize(), PAGE_SIZE);
    printf("env block size: %llx\n", envSize);  // 添加调试输出
    uc_err envErr = uc_mem_map(m_ucEngine, m_envBlockBase, envSize,
        UC_PROT_READ | UC_PROT_WRITE);
    if (envErr != UC_ERR_OK) {
        throw std::runtime_error("Failed to map environment block");
    }

    auto envData = this->GetEnvString();
    envErr = uc_mem_write(m_ucEngine, m_envBlockBase, envData.data(),
        envData.size() * sizeof(wchar_t));
    if (envErr != UC_ERR_OK) {
        throw std::runtime_error("Failed to write environment block");
    }

    for (DWORD i = 0; i < 64; i++) {
        GetTeb64()->TlsSlots[i] = (void*)0x1337ffffff;
    }
    for (DWORD i = 0; i < 64; i++) {
        GetTeb32()->TlsSlots[i] = 0x1337;
    }

}