#pragma once

#include "..\commons.h"

struct ProcessInfo
{
    PROCESSENTRY32W process;
    std::wstring imageName;
    std::wstring commandLine;
    std::wstring username;
    std::wstring integrity;
    bool canOpenProcess;
    bool alertableThread;
    bool propagateEnabled;
};

//
// Borrowed from modexp:
//  https://modexp.wordpress.com/2018/08/23/process-injection-propagate/
//  https://github.com/odzhan/injection/blob/master/propagate/propagate.c
//
typedef struct _win_props_t {
    DWORD  dwPid;
    WCHAR  ImageName[MAX_PATH];
    HANDLE hProperty;
    HWND   hParentWnd;
    HWND   hChildWnd;
    WCHAR  ParentClassName[MAX_PATH];
    WCHAR  ChildClassName[MAX_PATH];
} WINPROPS, *PWINPROPS;


struct PROPagateContext {
    std::vector<WINPROPS> windows;
    int maxName = 16;
    int maxClass = 16;
    bool bAll = true;
};

class RemoteInjectionFailedException
{
protected:
    std::wstring error_message;      ///< Error message

public:

    /** Constructor (C++ STL string, int, int).
     *  @param msg The error message
     */
    explicit
        RemoteInjectionFailedException(const std::wstring& msg) :
        error_message(msg)
    {}

    /** Destructor.
     *  Virtual to allow for subclassing.
     */
    virtual ~RemoteInjectionFailedException() throw () {}

    /** Returns a pointer to the (constant) error description.
     *  @return A pointer to a const wchar_t*. The underlying memory
     *  is in possession of the RemoteInjectionFailedException object. Callers must
     *  not attempt to free the memory.
     */
    virtual const wchar_t* what() const throw () {
        return error_message.c_str();
    }
};


extern bool globalVerboseOption;
extern bool globalAntiSplicingOption;


enum AllocationTechnique;
enum WriteTechnique;
enum ExecutionTechnique;



//
// Author: Elmue (https://stackoverflow.com/users/1487529)
// Source: 
//   https://stackoverflow.com/a/22949726
//
// Based on the 32 bit code of Sven B. Schreiber on:
// http://www.informit.com/articles/article.aspx?p=22442&seqNum=5
//
class NativeProcessInfo
{
public:
    NativeProcessInfo()
    {
#ifndef _WIN64
        assert(sizeof(PolonDefs::SYSTEM_THREAD) == 0x40 && sizeof(PolonDefs::SYSTEM_PROCESS) == 0xB8);
#else
        assert(sizeof(PolonDefs::SYSTEM_THREAD) == 0x50 && sizeof(PolonDefs::SYSTEM_PROCESS) == 0x100);
#endif

        mu32_DataSize = 1000;
        mp_Data = NULL;
    }
    virtual ~NativeProcessInfo()
    {
        if (mp_Data) LocalFree(mp_Data);
        mp_Data = NULL;
    }

    // Capture all running processes and all their threads.
    // returns an API or NTSTATUS Error code or zero if successfull
    DWORD Capture()
    {
        RESOLVE(ntdll, NtQuerySystemInformation);

        // This must run in a loop because in the mean time a new process may have started 
        // and we need more buffer than u32_Needed !!
        while (true)
        {
            if (!mp_Data)
            {
                mp_Data = reinterpret_cast<BYTE*>(LocalAlloc(LMEM_FIXED, mu32_DataSize));
                if (!mp_Data)
                    return GetLastError();
            }

            ULONG u32_Needed = 0;
            NTSTATUS s32_Status = _NtQuerySystemInformation(
                static_cast<SYSTEM_INFORMATION_CLASS>(PolonDefs::SystemProcessInformation),
                mp_Data, 
                mu32_DataSize, 
                &u32_Needed
            );

            if (s32_Status == STATUS_INFO_LENGTH_MISMATCH) // The buffer was too small
            {
                mu32_DataSize = u32_Needed + 4000;
                LocalFree(mp_Data);
                mp_Data = NULL;
                continue;
            }
            return s32_Status;
        }
    }

    // Searches a process by a given Process Identifier
    // Capture() must have been called before!
    PolonDefs::SYSTEM_PROCESS* FindProcessByPid(DWORD u32_PID)
    {
        if (!mp_Data)
        {
            assert(mp_Data);
            return NULL;
        }

        PolonDefs::SYSTEM_PROCESS* pk_Proc = (PolonDefs::SYSTEM_PROCESS*)mp_Data;
        while (TRUE)
        {
            if ((DWORD)(DWORD_PTR)pk_Proc->UniqueProcessId == u32_PID)
            {
                return pk_Proc;
            }

            if (!pk_Proc->NextEntryOffset)
            {
                return NULL;
            }

            pk_Proc = (PolonDefs::SYSTEM_PROCESS*)((BYTE*)pk_Proc + pk_Proc->NextEntryOffset);
        }
    }

    PolonDefs::SYSTEM_THREAD* FindThreadByTid(PolonDefs::SYSTEM_PROCESS* pk_Proc, DWORD u32_TID)
    {
        if (!pk_Proc)
        {
            assert(pk_Proc);
            return NULL;
        }

        // The first SYSTEM_THREAD structure comes immediately after the SYSTEM_PROCESS structure
        PolonDefs::SYSTEM_THREAD* pk_Thread = (PolonDefs::SYSTEM_THREAD*)((BYTE*)pk_Proc + sizeof(PolonDefs::SYSTEM_PROCESS));

        for (DWORD i = 0; i < pk_Proc->ThreadCount; i++)
        {
            if (pk_Thread->ClientID.UniqueThread == (HANDLE)(DWORD_PTR)u32_TID)
            {
                return pk_Thread;
            }

            pk_Thread++;
        }
        return NULL;
    }

    DWORD IsThreadSuspended(PolonDefs::SYSTEM_THREAD* pk_Thread, BOOL* pb_Suspended)
    {
        if (!pk_Thread)
        {
            return ERROR_INVALID_PARAMETER;
        }

        *pb_Suspended = (pk_Thread->ThreadState == PolonDefs::Waiting &&
            pk_Thread->WaitReason == PolonDefs::Suspended);

        return 0;
    }

private:
    BYTE*  mp_Data;
    DWORD  mu32_DataSize;
};

uint64_t milliseconds_now();

bool
IsProcessWow64(
    HANDLE process,
    bool* bIsWow64
);

bool IsProcessWow64(
    HANDLE process,
    BOOL* bIsWow64
);

uint32_t rotr32(
    uint32_t n, 
    unsigned int c
);

/// If delayInMiliseconds is 0 - set's that priority permanantely
void temporarilyBoostThreadPriority(
    HANDLE thread, 
    DWORD delayInMiliseconds
);

bool checkCrossArchRequirements(
    AllocationTechnique allocTech,
    WriteTechnique writeTech,
    ExecutionTechnique execTech,
    PROCESS_INFORMATION *pinfo
);

uint32_t rotl32(
    uint32_t n, 
    unsigned int c
);

BOOL getProcessModulePath(
    DWORD dwPid, 
    const BYTE* address, 
    LPWSTR ImageName, 
    DWORD dwSize,
    bool nameInsteadOfPath = false
);

std::wstring getModuleName(
    DWORD pid, 
    const BYTE* address
);

std::wstring getModulePath(
    DWORD pid, 
    const BYTE* address
);

std::wstring getProcessName(
    DWORD pid
);

const std::wstring getChildProcessPath(
    const std::wstring &targetProcess
);

void readFromPipe(
    PROCESS_INFORMATION piProcInfo, 
    HANDLE stdoutReadHandle
);

std::wstring parsePageProtection(
    DWORD value
);

DWORD interpretPageProtectionParameter(
    const std::wstring& param
);

bool checkIsThreadSuspended(
    DWORD pid,
    DWORD tid,
    bool *suspended
);

BOOL getProcessImageName(
    DWORD dwPid, 
    LPWSTR ImageName, 
    DWORD dwSize
);

std::pair<DWORD, HANDLE> getProcessMainThread(
    DWORD dwProcID
);

PROCESS_INFORMATION openRemoteProcess(
    DWORD dwPID,
    DWORD dwTid = 0,
    bool tryAlertableThread = false,
    bool closeHandles = false,
    DWORD desiredAccess = PROCESS_ALL_ACCESS
);

const std::wstring getProcessIntegrity(
    uint32_t processId
);

/*
std::vector<DWORD> getProcessThreads(
    DWORD dwProcID
);
*/

void printPageProtection(
    DWORD value
);

const std::wstring mapPageProtectionFlagValueToName(
    DWORD val
);

DWORD mapPageProtectionFlagNameToValue(
    const std::wstring& name
);

bool setupProcessMitigationPolicies();

BOOL setPrivilege(
    LPCTSTR lpszPrivilege,
    BOOL bEnablePrivilege,
    HANDLE mainToken = nullptr
);

DWORD interpretPageProtectionParameter(
    const std::wstring& param
);

bool antiSplicing(
    bool doIt
);

DWORD findExplorerPid(
    bool canBeGreater = false
);

void listRunningProcesses(
    bool detailed
);

uint32_t getProcessId(
    const std::wstring &processName,
    bool sameIntegrityLevel = false,
    bool returnParentPid = false
);

DWORD findSuitableProcess(
    bool canOpenProcess,
    std::vector<DWORD>* constrainedSetOfPids = nullptr,
    bool alertableThread = false
);

HANDLE MyOpenThread(
    DWORD dwDesiredAccess,
    BOOL  bInheritHandle,
    DWORD dwThreadId
);

void* allocate(
    uint8_t* buf,
    size_t size,
    DWORD protection
);

bool changeProtection(
    void* address,
    size_t size,
    DWORD newProtect,
    HANDLE process = nullptr,
    DWORD* poldProtect = nullptr
);

BOOL IsEntry(
    PWINPROPS e
);

BOOL CALLBACK propEnumProc(
    HWND hwnd, 
    LPWSTR lpszString, 
    HANDLE hData, 
    ULONG_PTR lParam
);

BOOL CALLBACK enumChildProc(
    HWND hwnd,
    LPARAM lParam
);

BOOL CALLBACK enumWindowsProc(
    HWND hwnd, 
    LPARAM lParam
);

bool sortEntries(
    const WINPROPS &a, 
    const WINPROPS &b
);

size_t getPropList(void);

std::vector<_win_props_t> enumerateWindowProperties(
    DWORD pid
);

std::vector<MODULEENTRY32W> collectProcessModules(
    DWORD dwPid
);

std::vector<MEMORY_BASIC_INFORMATION> collectProcessMemoryMap(
    HANDLE process
);

template<typename Filter>
void* findRemoteAllocation(const std::vector<MEMORY_BASIC_INFORMATION>& memoryMap, Filter filter)
{
    for (const auto allocation : memoryMap)
    {
        if (filter(allocation))
        {
            return allocation.BaseAddress;
        }
    }

    return nullptr;
}

template<typename Filter>
MODULEENTRY32W findProcessModule(
    const std::vector<MODULEENTRY32W>& modules, 
    Filter filter
)
{
    for (const auto module : modules)
    {
        if (filter(module))
        {
            return module;
        }
    }

    return {};
}

template<typename Filter>
MEMORY_BASIC_INFORMATION findRemoteAllocation2(
    const std::vector<MEMORY_BASIC_INFORMATION>& memoryMap, 
    Filter filter
)
{
    for (const auto allocation : memoryMap)
    {
        if (filter(allocation))
        {
            return allocation;
        }
    }

    return {};
}

bool downloadResource(
    const std::wstring& url,
    std::vector<uint8_t> &buffer,
    bool* dontRetry = nullptr
);

bool _downloadResource(
    const std::wstring& url,
    std::vector<uint8_t> &buffer,
    bool* dontRetry = nullptr
);

void unloadModule(
    const wchar_t* module
);

void findOwnAllocationMemoryRegion(
    size_t& moduleSize, 
    uintptr_t& alloc, 
    bool& isMapped
);

void _findOwnAllocationMemoryRegion(
    size_t& moduleSize, 
    uintptr_t& alloc, 
    bool& isMapped
);

bool getAllocationModulePE(
    const MEMORY_BASIC_INFORMATION& me, 
    LPBYTE& peBaseAddress, 
    size_t& moduleSizeFromPeHeader
);

bool analyseOwnMemoryRegion(
    PE& pe
);

bool analyseOwnMemoryRegion(
    size_t moduleSize, 
    uintptr_t alloc, 
    PE& pe, 
    bool isMapped
);

bool initialiseExceptionHandlers(
    uintptr_t baseAddress, 
    PE& pe
);

bool initialiseNotMappedModule();

std::vector<PROCESSENTRY32W> collectRunningProcesses();