#pragma once
#include "head.h"

#define PAGE_SIZE 0x1000
#define CF_MASK (1 << 0)
#define PF_MASK (1 << 2)
#define ZF_MASK (1 << 6)
#define SF_MASK (1 << 7)
#define OF_MASK (1 << 11)
#define ALL_MASK (OF_MASK | SF_MASK | ZF_MASK | PF_MASK | CF_MASK)
#define STACK_BASE_64 0x7ffffffde000
#define STACK_BASE_32 0xfffdd000
#define STACK_SIZE_64 0x40000
#define STACK_SIZE_32 0x21000
#define HEAP_ADDRESS_64 0x500000000
#define HEAP_SIZE_64 0x5000000
#define HEAP_ADDRESS_32 0x5000000
#define HEAP_SIZE_32 0x5000000

#define PEB_BASE 0x80000
#define TEB_BASE 0x90000

#define X86_GDT_ADDR 0x30000
#define X86_GDT_LIMIT 0x1000
#define X86_GDT_ENTRY_SIZE 0x8

#define API_FUNCTION_SIZE 8
#define PAGE_ALIGN(Va) (ULONG_PTR)(Va) & ~(PAGE_SIZE - 1)
#define PAGE_ALIGN_64(Va) (Va) & ~(0x1000ull - 1)
#define PAGE_ALIGN_64k(Va) ((Va)) & ~(0x10000ull - 1)
#define AlignSize(Size, Align) (Size + Align - 1) / Align* Align

enum class WinVer {
    kWin7 = 0x0610,
    kWin7SP1 = 0x0611,
    kWin8 = 0x0620,
    kWin81 = 0x0630,
    kWin10 = 0x0A00,
    kWin10RS1 = 0x0A01,   // Anniversary update
    kWin10RS2 = 0x0A02,   // Creators update
    kWin10RS3 = 0x0A03,   // Fall creators update
    kWin10RS4 = 0x0A04,   // Spring creators update
    kWin10RS5 = 0x0A05,   // October 2018 update
    kWin1019H1 = 0x0A06,  // May 2019 update 19H1
    kWin1019H2 = 0x0A07,  // November 2019 update 19H2
    kWin1020H1 = 0x0A08   // April 2020 update 20H1
};

class Sandbox {
   public:
    Sandbox();
    ~Sandbox();

    // Public methods
    auto InitEnv(std::shared_ptr<BasicPeInfo> peInfo) -> void;

    auto Run() -> void;

    auto GetCapstoneHandle() const -> csh { return m_csHandle; }
    auto GetUnicornHandle() const -> uc_engine* { return m_ucEngine; }
    auto GetPeInfo() const -> std::shared_ptr<BasicPeInfo> { return m_peInfo; }

   private:
    std::shared_ptr<BasicPeInfo> m_peInfo;
    uint64_t m_gsBase;
    uint64_t m_pebBase;
    uint64_t m_pebEnd;
    uint64_t m_tebBase;
    uint64_t m_tebEnd;
    PVOID m_stackBuffer;  // 没有释放
    uint64_t m_stackBase;
    uint64_t m_stackSize;
    uint64_t m_stackEnd;
    uint64_t m_heapBase;
    uint64_t m_heapSize;
    uint64_t m_heapEnd;
    uint64_t m_fakeBase;
    struct_gs_base m_gsBaseStruct;
    X64TEB m_teb64;
    X64PEB m_peb64;
    X32TEB m_teb32;
    X32PEB m_peb32;
    csh m_csHandle;         // Capstone handle
    uc_engine* m_ucEngine;  // Unicorn engine handle
    std::vector<std::shared_ptr<moudle_import>> m_impFuncDict;
    std::vector<std::shared_ptr<moudle_export>> m_exportFuncDict;
    std::vector<std::shared_ptr<struct_moudle>> m_moduleList;
    auto ResoveImport() -> void;
    auto SetupVirtualMachine() -> void;
    auto ResolveExport() -> void;
    auto PushModuleToVM(const char* dllName, uint64_t moduleBase,
                        uint32_t x32Base) -> void;
    auto processImportModule(const moudle_import* importModule) -> void;
};