#ifdef _MSC_VER
#include <WinSock2.h>
#include <WS2tcpip.h>
#include <iphlpapi.h>
#include <sddl.h>
#endif

#include "SocketOwnership.hpp"
#include <vector>
#include <algorithm>

#ifdef _MSC_VER

std::optional<std::string> getStringSidForProcess(HANDLE processHandle) {
    HANDLE tokenHandle = NULL;
    if (!OpenProcessToken(processHandle, TOKEN_READ, &tokenHandle) || !tokenHandle) {
        CloseHandle(processHandle);
        return {};
    }

    DWORD tokenUserLength = 0;
    GetTokenInformation(tokenHandle, TokenUser, NULL, 0, &tokenUserLength);
    if (tokenUserLength == 0) {
        CloseHandle(tokenHandle);
        CloseHandle(processHandle);
        return {};
    }

    PTOKEN_USER tokenUser = (PTOKEN_USER)malloc((size_t)tokenUserLength);
    if (!GetTokenInformation(tokenHandle, TokenUser, tokenUser, tokenUserLength, &tokenUserLength)) {
        CloseHandle(tokenHandle);
        CloseHandle(processHandle);
        return {};
    }

    LPTSTR stringSid = NULL;
    ConvertSidToStringSid(tokenUser->User.Sid, &stringSid);
    std::optional<std::string> ssid = {};
    if (stringSid) {
        ssid = std::string(stringSid);
        LocalFree(stringSid);
    }

    free(tokenUser);
    CloseHandle(tokenHandle);
    CloseHandle(processHandle);
    return ssid;
}

std::optional<std::string> findLocalhostRemoteSocketOwner(int port) {
    DWORD tcpTableSize = 1024;
    DWORD wantTcpTableSize = tcpTableSize;
    PMIB_TCPTABLE2 tcpTable = (PMIB_TCPTABLE2)malloc((size_t)tcpTableSize);

    ULONG res = 0;

    while ((res = GetTcpTable2(tcpTable, &wantTcpTableSize, TRUE)) != NO_ERROR) {
        if (res != ERROR_INSUFFICIENT_BUFFER) {
            free(tcpTable);
            return {};
        }
        wantTcpTableSize = tcpTableSize = (DWORD)std::max((size_t)wantTcpTableSize + 16, (size_t)tcpTableSize);
        free(tcpTable);
        tcpTable = (PMIB_TCPTABLE2)malloc((size_t)tcpTableSize);
        if (tcpTable == NULL) {
            return {};
        }
    }

    DWORD owningPid = -1;
    for (DWORD i = 0; i < tcpTable->dwNumEntries; i++) {
        auto& row = tcpTable->table[i];
        if (row.dwLocalAddr == htonl(INADDR_LOOPBACK) && ntohs(row.dwLocalPort) == port) {
            owningPid = row.dwOwningPid;
            break;
        }
    }
    free(tcpTable);

    if (owningPid == -1) {
        return {};
    }
     
    HANDLE processHandle = OpenProcess(PROCESS_QUERY_LIMITED_INFORMATION, FALSE, owningPid);
    if (processHandle == NULL) {
        return {};
    }

    return getStringSidForProcess(processHandle);
}

#else /* _MSC_VER */

#include <string>
#include <libproc.h>
#include <netinet/in.h>
#include <netinet/tcp.h>
#include <arpa/inet.h>

#define IPv6_2_IPv4(v6) (((uint8_t *)((struct in6_addr *)v6)->s6_addr) + 12)

std::optional<uid_t> findLocalhostRemoteSocketOwner(int port) {
    int nb = 0;
    std::vector<int> pids;
    
    for(int i = 0; i < 100; i++) {
        if((nb = proc_listpids(PROC_ALL_PIDS, 0, pids.data(), (int)(pids.size() * sizeof(int)))) <= 0) {
            return {};
        }
        if(nb <= pids.size() * sizeof(int)) {
            pids.resize(nb / sizeof(int));
            break;
        }
        pids.resize(std::max(pids.size(), nb / sizeof(int) + 5));
    }
    
    for(auto p : pids) {
        if(!p) {
            continue;
        }
        struct proc_taskallinfo tai;
    
        nb = proc_pidinfo(p, PROC_PIDTASKALLINFO, 0, &tai, sizeof(tai));
        if (nb <= 0 || nb < sizeof(tai)) {
            continue;
        }
        
        std::vector<struct proc_fdinfo> infos;
        infos.resize(tai.pbsd.pbi_nfiles + 20);
        if((nb = proc_pidinfo(p, PROC_PIDLISTFDS, 0, infos.data(), (int)(infos.size() * sizeof(struct proc_fdinfo)))) < 0) {
            continue;
        }
        infos.resize(nb / sizeof(struct proc_fdinfo));
        
        if(infos.empty()) {
            continue;
        }

        for(auto & info : infos) {
            if(info.proc_fdtype != PROX_FDTYPE_SOCKET) {
                continue;
            }
            
            struct socket_fdinfo si;
            nb = proc_pidfdinfo(p, info.proc_fd, PROC_PIDFDSOCKETINFO, &si, sizeof(si));
            
            if(nb < sizeof(si)) {
                continue;
            }
            
            if(si.psi.soi_kind != SOCKINFO_TCP) {
                continue;
            }
            
            std::string localAddress;
            int localPort = 0;
            switch(si.psi.soi_family) {
                case AF_INET:
                {
                    char buf[INET_ADDRSTRLEN] = {0};
                    inet_ntop(AF_INET, &si.psi.soi_proto.pri_tcp.tcpsi_ini.insi_laddr.ina_46.i46a_addr4, buf, INET_ADDRSTRLEN);
                    localAddress = buf;
                    localPort = (int)ntohs(si.psi.soi_proto.pri_tcp.tcpsi_ini.insi_lport);
                    break;
                }
                case AF_INET6:
                {
                    unsigned char * addr = (unsigned char *)&si.psi.soi_proto.pri_tcp.tcpsi_ini.insi_laddr.ina_6;
                    localPort = (int)ntohs(si.psi.soi_proto.pri_tcp.tcpsi_ini.insi_lport);
                    if((si.psi.soi_proto.pri_tcp.tcpsi_ini.insi_vflag & INI_IPV4) != 0) {
                        char buf[INET_ADDRSTRLEN] = {0};
                        inet_ntop(AF_INET, IPv6_2_IPv4(addr), buf, INET_ADDRSTRLEN);
                        localAddress = buf;
                    } else {
                        char buf[INET6_ADDRSTRLEN] = {0};
                        inet_ntop(AF_INET6, addr, buf, INET6_ADDRSTRLEN);
                        localAddress = buf;
                    }
                    break;
                }
            }
            
            if(localPort == 0 || localAddress != "127.0.0.1") {
                continue;
            }
            
            if(localPort == port) {
                return tai.pbsd.pbi_uid;
            }
        }
    }
    
    return {};
}

#endif /* _MSC_VER */