#include "pch.h"
#include "WinDivert/include/windivert.h"

#include <strsafe.h>
#include <spdlog/spdlog.h>
#include "globalcfg.h"
#include "usrerr.h"

#pragma comment(lib, "WinDivert.lib")

typedef struct {
    bool   isRunning;
    HANDLE proxyThread;
    HANDLE proxyHandle;
    bool   isExitSvr;
} UDP_PROXY_CTX, *PUDP_PROXY_CTX;

typedef struct {
    TCHAR         streamFilter[1024];
    TCHAR         targetIp[MAX_IP_LEN];
    UINT16        targetPort;
    int           vmId;
    int           svrId;
    UDP_PROXY_CTX udpCtx;
} PROXY_INFO, *PPROXY_INFO;


static std::unordered_map<std::string, PPROXY_INFO> g_ProxyColleagues;


int CreatePorxyService(const TCHAR *pTargetIp, int targetPort, int vmId, int svrId) {
    //int CreatePorxyService() {
    //int                                                    ret;
    static HANDLE                                          hDriver;
    PPROXY_INFO                                            pInfo;
    std::string                                            key;
    std::unordered_map<std::string, PPROXY_INFO>::iterator iter;

    // 查找先前代理是否存在
    key = std::string(pTargetIp) + ":" + std::to_string(targetPort);

    if ((iter = g_ProxyColleagues.find(key)) != g_ProxyColleagues.end()) {
        pInfo = iter->second;
        // 如果配置完全相同则直接返回
        if (pInfo->vmId == vmId && pInfo->svrId == svrId) {
            return ERR_SUCCESS;
        } else {
            pInfo->vmId    = vmId;
            pInfo->svrId   = svrId;
        }
    } else {
        // 创建新的代理配置
        pInfo = static_cast<PPROXY_INFO>(HeapAlloc(GetProcessHeap(), 0, sizeof(PROXY_INFO)));

        if (pInfo == nullptr) {
            SPDLOG_ERROR(TEXT("Error allocating {0} bytes memory "), sizeof(PROXY_INFO));
            return -ERR_MALLOC_MEMORY;
        }

        memset(pInfo, 0, sizeof(PROXY_INFO));
    }

    memset(pInfo->streamFilter, 0, 1024);

    StringCbPrintf(pInfo->streamFilter, 1024, TEXT("(udp.DstPort == %d and )"), targetPort);

    hDriver = WinDivertOpen(pInfo->streamFilter, WINDIVERT_LAYER_NETWORK, 0, 0);

    if (hDriver == INVALID_HANDLE_VALUE) {
        //ERROR_INSUFFICIENT_BUFFER
        SPDLOG_ERROR(TEXT("Open Driver With Filter \"{1}\" Error: {0}"), GetLastError(), pInfo->streamFilter);
        return -ERR_SYS_CALL;
    }

    CreateThread(
        nullptr,    // Thread attributes
        0,          // Stack size (0 = use default)
        [](LPVOID lpParameter) {
            const HANDLE      h = lpParameter;
            unsigned char     packet[2048];
            UINT              packet_len;
            WINDIVERT_ADDRESS addr;

            while (true) {
                if (!WinDivertRecv(h, packet, 0xFFFF, &packet_len, &addr)) {
                    SPDLOG_ERROR(TEXT("failed to read packet ({0})"), GetLastError());
                    continue;
                }

                switch (addr.Event) {
                    case WINDIVERT_EVENT_SOCKET_BIND:
                        SPDLOG_ERROR("BIND");
                        break;
                    case WINDIVERT_EVENT_SOCKET_LISTEN:

                        SPDLOG_ERROR("LISTEN");
                        break;
                    case WINDIVERT_EVENT_SOCKET_CONNECT:

                        SPDLOG_ERROR("CONNECT: {0:x}, {1:x}", addr.Socket.EndpointId, addr.Socket.ParentEndpointId);
                        break;
                    case WINDIVERT_EVENT_SOCKET_ACCEPT:

                        SPDLOG_ERROR("ACCEPT");
                        break;
                    case WINDIVERT_EVENT_SOCKET_CLOSE:

                        SPDLOG_ERROR("CLOSE: {0:x}, {1:x}", addr.Socket.EndpointId, addr.Socket.ParentEndpointId);
                        break;
                    default:
                        SPDLOG_ERROR(TEXT("***({0})"), static_cast<int>(addr.Event));
                        break;
                }

                /*if (!WinDivertSendEx(h, packet, packet_len, nullptr, 0, &addr, sizeof(WINDIVERT_ADDRESS), nullptr)) {
                    SPDLOG_ERROR(TEXT("warning: failed to reinject packet ({0})\n"), GetLastError());
                }*/
            }

            return static_cast<DWORD>(0);
        },           // Thread start address
        hDriver,     // Parameter to pass to the thread
        0,           // Creation flags
        nullptr);    // Thread id
    return 0;
}