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

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

#include <winsock2.h>

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

typedef enum {
    PROXY_TCP = 1 << 0,
    PROXY_UDP = 1 << 1,
} PROXY_PORTO_TYPE;

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

typedef struct {
    SOCKET proxySock;
    SOCKET reservePortSock;
    SOCKET altSock;
    bool   isRunning;
    UINT16 proxyPort;
    UINT16 altPort;
    HANDLE proxyThread;
    HANDLE proxyHandle;
    bool   isExitSvr;
} TCP_PROXY_CTX, *PTCP_PORXY_CTX;

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

typedef struct {
    SOCKET      s;
    PPROXY_INFO pCtx;
    bool        inbound;
    in_addr     dest;
} PROXY_CONNECTION_CONFIG, *PPROXY_CONNECTION_CONFIG;

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

static int NewAvaliableSocket(SOCKET *pSock, UINT16 *pPort) {
    const int    reusedAddr = 1;
    sockaddr_in  proxySvrAddr {};
    sockaddr_in  bindAddr {};
    int          ret;
    int          addrSize = sizeof(sockaddr_in);
    const SOCKET sock     = socket(AF_INET, SOCK_STREAM, 0);

    if (sock == INVALID_SOCKET) {
        SPDLOG_ERROR(TEXT("Cretate TCP Socket error: {0}"), WSAGetLastError());
        return -ERR_SOCKET_BIND_PORT;
    }

    if (setsockopt(sock, SOL_SOCKET, SO_REUSEADDR, reinterpret_cast<const char *>(&reusedAddr), sizeof(int)) ==
        SOCKET_ERROR) {
        SPDLOG_ERROR(TEXT("Failed to set socket optino SO_REUSEADDR ({0}})"), WSAGetLastError());
        closesocket(sock);
        return -ERR_SOCKET_SET_OPT;
    }

    proxySvrAddr.sin_family      = AF_INET;
    proxySvrAddr.sin_addr.s_addr = htonl(INADDR_ANY);
    proxySvrAddr.sin_port        = htons(0);

    if (bind(sock, reinterpret_cast<SOCKADDR *>(&proxySvrAddr), sizeof(proxySvrAddr)) == SOCKET_ERROR) {
        SPDLOG_ERROR(TEXT("Failed to bind socket ({0})"), WSAGetLastError());
        closesocket(sock);
        return -ERR_SOCKET_BIND;
    }

    if (getsockname(sock, reinterpret_cast<struct sockaddr *>(&bindAddr), &addrSize) != 0) {
        SPDLOG_ERROR(TEXT("Failed to get socket bind port ({0})"), WSAGetLastError());
        closesocket(sock);
        return -ERR_SOCKET_GET_OPT;
    }

    *pSock = sock;
    *pPort = ntohs(bindAddr.sin_port);

    return ERR_SUCCESS;
}

static DWORD proxyForwardCb(LPVOID lpParameter) {
    return 0;
}

static DWORD proxyConnCb(LPVOID lpParameter) {
    const auto               p = static_cast<PPROXY_CONNECTION_CONFIG>(lpParameter);
    sockaddr_in              altAddr {};
    SOCKET                   altSock;
    HANDLE                   thread;
    SOCKET                   connSock = p->s;
    in_addr                  dest     = p->dest;
    PPROXY_INFO              pInfo    = p->pCtx;
    PPROXY_CONNECTION_CONFIG pCfgProxySvr, pCfgAltSvr;

    HeapFree(GetProcessHeap(), 0, p);

    altSock = socket(AF_INET, SOCK_STREAM, 0);

    if (altSock == INVALID_SOCKET) {
        SPDLOG_ERROR(TEXT("Cretate TCP Socket error: {0}"), WSAGetLastError());
        closesocket(connSock);
        return 0;
    }

    memset(&altAddr, 0, sizeof(altAddr));
    altAddr.sin_family = AF_INET;
    altAddr.sin_port   = htons(pInfo->tcpCtx.altPort);
    altAddr.sin_addr   = dest;

    if (connect(altSock, reinterpret_cast<SOCKADDR *>(&altAddr), sizeof(altAddr)) == SOCKET_ERROR) {
        SPDLOG_ERROR(TEXT("Failed to connect socket ({0})"), WSAGetLastError());
        closesocket(connSock);
        closesocket(altSock);
        return 0;
    }

    pInfo->tcpCtx.altSock = altSock;

    pCfgProxySvr = static_cast<PPROXY_CONNECTION_CONFIG>(
        HeapAlloc(GetProcessHeap(), 0, sizeof(PROXY_CONNECTION_CONFIG)));

    if (pCfgProxySvr == nullptr) {
        SPDLOG_ERROR(TEXT("Malloc {0} bytes error."), sizeof(PROXY_CONNECTION_CONFIG));
        closesocket(connSock);
        closesocket(altSock);
        return 0;
    }

    pCfgAltSvr = static_cast<PPROXY_CONNECTION_CONFIG>(HeapAlloc(GetProcessHeap(), 0, sizeof(PROXY_CONNECTION_CONFIG)));

    if (pCfgAltSvr == nullptr) {
        SPDLOG_ERROR(TEXT("Malloc {0} bytes error."), sizeof(PROXY_CONNECTION_CONFIG));
        closesocket(connSock);
        closesocket(altSock);
        free(pCfgProxySvr);
        return 0;
    }

    thread = CreateThread(nullptr,           // Thread attributes
                          0,                 // Stack size (0 = use default)
                          proxyForwardCb,    // Thread start address
                          pCfgProxySvr,      // Parameter to pass to the thread
                          0,                 // Creation flags
                          nullptr);          // Thread id

    if (thread == nullptr) {
        SPDLOG_ERROR(TEXT("Failed to create thread ({0})"), GetLastError());
        closesocket(connSock);
        closesocket(altSock);
        free(pCfgProxySvr);
        free(pCfgAltSvr);
        return 0;
    }

    return 0;
}

static int StartTcpProxyService(PPROXY_INFO pInfo) {
    pInfo->tcpCtx.isExitSvr = false;

    if (listen(pInfo->tcpCtx.proxySock, SOMAXCONN) == SOCKET_ERROR) {
        SPDLOG_ERROR(TEXT("Failed to listen socket ({0})"), WSAGetLastError());
        closesocket(pInfo->tcpCtx.proxySock);
        return -ERR_SOCKET_LISTEN;
    }

    pInfo->tcpCtx.proxyThread = CreateThread(
        nullptr,    // Thread attributes
        0,          // Stack size (0 = use default)
        [](LPVOID lpParameter) {
            const auto p = static_cast<PPROXY_INFO>(lpParameter);

            while (!p->tcpCtx.isExitSvr) {
                PPROXY_CONNECTION_CONFIG pCfg;
                sockaddr_in              connAddr {};
                int                      size = sizeof(connAddr);

                SOCKET s = accept(p->tcpCtx.proxySock, reinterpret_cast<struct sockaddr *>(&connAddr), &size);

                if (s == INVALID_SOCKET) {
                    continue;
                }

                pCfg = static_cast<PPROXY_CONNECTION_CONFIG>(
                    HeapAlloc(GetProcessHeap(), 0, sizeof(PROXY_CONNECTION_CONFIG)));

                if (pCfg) {
                    HANDLE thread;
                    pCfg->s    = s;
                    pCfg->dest = connAddr.sin_addr;
                    pCfg->pCtx = p;

                    thread = CreateThread(nullptr,
                                          0,
                                          proxyConnCb,    // Thread start address
                                          pCfg,           // Parameter to pass to the thread
                                          0,
                                          nullptr);

                    if (thread == nullptr) {
                        closesocket(s);
                        HeapFree(GetProcessHeap(), 0, pCfg);
                        continue;
                    }

                    CloseHandle(thread);
                }
            }

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

    if (pInfo->tcpCtx.proxyThread == nullptr) {
        SPDLOG_ERROR(TEXT("Create TCP Listen Thread Error ({0})"), GetLastError());
        return -ERR_CREATE_THREAD;
    }
}

int CreatePorxyService(int proType, 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 && pInfo->proType == proType) {
            return ERR_SUCCESS;
        } else {
            pInfo->vmId    = vmId;
            pInfo->svrId   = svrId;
            pInfo->proType = proType;
        }
    } 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);

    // 重构过滤器
    if (pInfo->proType & PROXY_UDP) {
        StringCbPrintf(pInfo->streamFilter, 1024, TEXT("(udp.DstPort == %d)"), targetPort);
        //StringCbPrintf(pFilter, 1024, TEXT("tcp && remotePort == 9276"));
    }

    if (pInfo->proType & PROXY_TCP) {
        TCHAR tmpFilter[1024];

        // 预先分配代理和转发 Socket
        if ((ret = NewAvaliableSocket(&pInfo->tcpCtx.proxySock, &pInfo->tcpCtx.proxyPort)) != ERR_SUCCESS) {
            return ret;
        }

        if ((ret = NewAvaliableSocket(&pInfo->tcpCtx.reservePortSock, &pInfo->tcpCtx.altPort)) != ERR_SUCCESS) {
            return ret;
        }

        // 构建过滤器
        StringCbPrintf(tmpFilter,
                       1024,
                       TEXT("(tcp.DstPort == %d or tcp.DstPort == %d or tcp.DstPort == %d or tcp.SrcPort == %d or "
                            "tcp.SrcPort == %d or tcp.SrcPort == %d)"),
                       targetPort,
                       pInfo->tcpCtx.proxyPort,
                       pInfo->tcpCtx.altPort,
                       targetPort,
                       pInfo->tcpCtx.proxyPort,
                       pInfo->tcpCtx.altPort);

        if (lstrlen(pInfo->streamFilter) > 0) {
            StringCbCat(pInfo->streamFilter, 1024, TEXT(" or "));
            StringCbCat(pInfo->streamFilter, 1024, tmpFilter);
        } else {
            StringCbCopy(pInfo->streamFilter, 1024, tmpFilter);
        }
    }

    // 启动代理服务
    if ((ret = StartTcpProxyService(pInfo)) != ERR_SUCCESS) {
        return ret;
    }

    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[0xFFFF];
            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;
}