#include "pch.h"
#include "tunnel.h"
#include "usrerr.h"
#include <strsafe.h>
#include <tchar.h>
#include <Iphlpapi.h>
#include <iostream>

#include "globalcfg.h"
#include "misc.h"

#include <shlwapi.h>

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

static const TCHAR *g_TabCIDR[] = {
    TEXT("128.0.0.0"),       TEXT("192.0.0.0"),       TEXT("224.0.0.0"),       TEXT("240.0.0.0"),
    TEXT("248.0.0.0"),       TEXT("252.0.0.0"),       TEXT("254.0.0.0"),       TEXT("255.0.0.0"),
    TEXT("255.128.0.0"),     TEXT("255.192.0.0"),     TEXT("255.224.0.0"),     TEXT("255.240.0.0"),
    TEXT("255.248.0.0"),     TEXT("255.252.0.0"),     TEXT("255.254.0.0"),     TEXT("255.255.0.0"),
    TEXT("255.255.128.0"),   TEXT("255.255.192.0"),   TEXT("255.255.224.0"),   TEXT("255.255.240.0"),
    TEXT("255.255.248.0"),   TEXT("255.255.252.0"),   TEXT("255.255.254.0"),   TEXT("255.255.255.0"),
    TEXT("255.255.255.128"), TEXT("255.255.255.192"), TEXT("255.255.255.224"), TEXT("255.255.255.240"),
    TEXT("255.255.255.248"), TEXT("255.255.255.252"), TEXT("255.255.255.254"), TEXT("255.255.255.255")};

int NetmaskToCIDR(const TCHAR *pNetMask) {

    for (int i = 0; i < static_cast<int>(std::size(g_TabCIDR)); i++) {
        if (lstrcmp(g_TabCIDR[i], pNetMask) == 0) {
            return i + 1;
        }
    }

    return 0xFF;
}

const TCHAR *CIDRToNetmask(const UINT8 cidr) {

    if (cidr >= 1 && cidr <= std::size(g_TabCIDR)) {
        return g_TabCIDR[cidr - 1];
    }

    return TEXT("");
}

int GetWindowsHyperVStatus(int *pEnabled) {
    int   ret;
    TCHAR cmdBuf[MAX_PATH];
    TCHAR cmdResult[2048];
    DWORD retCode;

    if (pEnabled == nullptr) {
        SPDLOG_ERROR("Input pEnabled params error");
        return -ERR_INPUT_PARAMS;
    }

    if (FAILED(StringCbPrintf(cmdBuf,
                              MAX_PATH,
                              "PowerShell -Command \"& {Get-WindowsOptionalFeature -FeatureName Microsoft-Hyper-V-All "
                              "-Online | Format-List -Property State}\""))) {
        SPDLOG_ERROR("Format String Error");
        return -ERR_MEMORY_STR;
    }

    if ((ret = RunCommand(cmdBuf, cmdResult, 2048, &retCode)) != ERR_SUCCESS) {
        SPDLOG_ERROR("Run command [{0}] error: {1}", cmdBuf, ret);
        return -ERR_CALL_SHELL;
    }

    if (StrStr(cmdResult, TEXT("Enabled")) != nullptr) {
        *pEnabled = TRUE;
    } else {
        *pEnabled = FALSE;
    }

    SPDLOG_DEBUG("Run Get Windows Hyper-V status Command({1}): {0} result: {2}\n", cmdBuf, retCode, cmdResult);

    return ERR_SUCCESS;
}

int EnableWindowsHyperV(bool enabled) {
    int   ret;
    TCHAR cmdBuf[MAX_PATH];
    DWORD retCode;

    if (enabled) {
        if (FAILED(StringCbPrintf(cmdBuf,
                                  MAX_PATH,
                                  "PowerShell -Command \"& {Enable-WindowsOptionalFeature -Online -FeatureName "
                                  "Microsoft-Hyper-V -All}\""))) {
            SPDLOG_ERROR("Format String Error");
            return -ERR_MEMORY_STR;
        }
    } else {
        if (FAILED(StringCbPrintf(cmdBuf,
                                  MAX_PATH,
                                  "PowerShell -Command \"& {Disable-WindowsOptionalFeature -Online -FeatureName "
                                  "Microsoft-Hyper-V-Hypervisor}\""))) {
            SPDLOG_ERROR("Format String Error");
            return -ERR_MEMORY_STR;
        }
    }

    if ((ret = RunCommand(cmdBuf, nullptr, 0, &retCode)) != ERR_SUCCESS) {
        SPDLOG_ERROR("Run command [{0}] error: {1}", cmdBuf, ret);
        return -ERR_CALL_SHELL;
    }

    if (retCode != 0) {
        SPDLOG_ERROR("PowerShell return error({1}): {0}", cmdBuf, retCode);
        return -ERR_PROCESS_RETURN;
    }

    return ERR_SUCCESS;
}

int GetInterfaceIndexByName(const TCHAR *pInterfaceName, int *pIndex) {
    int   ret;
    DWORD retCode;
    TCHAR cmdBuf[MAX_PATH];
    TCHAR cmdResult[MAX_PATH] = {};

    if (pInterfaceName == nullptr || lstrlen(pInterfaceName) == 0) {
        SPDLOG_ERROR("Input pInterfaceName params error: {0}", pInterfaceName);
        return -ERR_INPUT_PARAMS;
    }

    if (pIndex == nullptr) {
        SPDLOG_ERROR("Input pIndex params error");
        return -ERR_INPUT_PARAMS;
    }

    if (FAILED(
            StringCbPrintf(cmdBuf,
                           MAX_PATH,
                           "PowerShell -Command \"& {Get-NetAdapter -Name %s | Format-List -Property InterfaceIndex}\"",
                           pInterfaceName))) {
        SPDLOG_ERROR("Format String Error");
        return -ERR_MEMORY_STR;
    }

    if ((ret = RunCommand(cmdBuf, cmdResult, MAX_PATH, &retCode)) != ERR_SUCCESS) {
        SPDLOG_ERROR("Run command [{0}] error: {1}", cmdBuf, ret);
        return -ERR_CALL_SHELL;
    }

    SPDLOG_DEBUG("Run command [{0}] resutl \'{1}\' return {2}", cmdBuf, cmdResult, retCode);

    // 移除
    StringRemoveAll(cmdResult, TEXT("\r\n"));
    StringRemoveAll(cmdResult, TEXT(" "));
    StringRemoveAll(cmdResult, TEXT("InterfaceIndex:"));

    *pIndex = strtol(cmdResult, nullptr, 10);

    return ERR_SUCCESS;
}

int SetNATRule(const TCHAR *pInterfaceName, const TCHAR *pCidrIpaddr) {
    int   ret;
    TCHAR cmdBuf[MAX_PATH];
    DWORD retCode;

    if (pInterfaceName == nullptr || lstrlen(pInterfaceName) == 0) {
        SPDLOG_ERROR("Input pInterfaceName params error: {0}", pInterfaceName);
        return -ERR_INPUT_PARAMS;
    }

    if (pCidrIpaddr == nullptr || lstrlen(pCidrIpaddr) == 0) {
        SPDLOG_ERROR("Input pCidrIpaddr params error: {0}", pCidrIpaddr);
        return -ERR_INPUT_PARAMS;
    }

    if (FAILED(
            StringCbPrintf(cmdBuf,
                           MAX_PATH,
                           "PowerShell -Command \"& {New-NetNat -Name %s_nat -InternalIPInterfaceAddressPrefix %s}\"",
                           pInterfaceName,
                           pCidrIpaddr))) {
        SPDLOG_ERROR("Format String Error");
        return -ERR_MEMORY_STR;
    }

    if ((ret = RunCommand(cmdBuf, nullptr, 0, &retCode)) != ERR_SUCCESS) {
        SPDLOG_ERROR("Run command [{0}] error: {1}", cmdBuf, ret);
        return -ERR_CALL_SHELL;
    }

    SPDLOG_DEBUG("Run Set IP Command({1}): {0}", cmdBuf, retCode);

    if (retCode != 0) {
        SPDLOG_ERROR("PowerShell return error({1}): {0}", cmdBuf, retCode);
        return -ERR_PROCESS_RETURN;
    }

    return ERR_SUCCESS;
}

int RemoveNATRule(const TCHAR *pInterfaceName) {
    int   ret;
    TCHAR cmdBuf[MAX_PATH];
    DWORD retCode;

    if (pInterfaceName == nullptr || lstrlen(pInterfaceName) == 0) {
        SPDLOG_ERROR("Input pInterfaceName params error: {0}", pInterfaceName);
        return -ERR_INPUT_PARAMS;
    }

    if (FAILED(StringCbPrintf(cmdBuf,
                              MAX_PATH,
                              "PowerShell -Command \"& {Remove-NetNat -Name %s_nat -Confirm:$false}\"",
                              pInterfaceName))) {
        SPDLOG_ERROR("Format String Error");
        return -ERR_MEMORY_STR;
    }

    if ((ret = RunCommand(cmdBuf, nullptr, 0, &retCode)) != ERR_SUCCESS) {
        SPDLOG_ERROR("Run command [{0}] error: {1}", cmdBuf, ret);
        return -ERR_CALL_SHELL;
    }

    SPDLOG_DEBUG("Run Set IP Command({1}): {0}", cmdBuf, retCode);

    return ERR_SUCCESS;
}

int RemoveInterfaceIpAddress(const TCHAR *pInterfaceName) {
    int   ret;
    TCHAR cmdBuf[MAX_PATH];
    DWORD retCode;

    if (pInterfaceName == nullptr || lstrlen(pInterfaceName) == 0) {
        SPDLOG_ERROR("Input pInterfaceName params error: {0}", pInterfaceName);
        return -ERR_INPUT_PARAMS;
    }

    if (FAILED(StringCbPrintf(cmdBuf,
                              MAX_PATH,
                              "PowerShell -Command \"& {Remove-NetIPAddress -InterfaceAlias %s  -Confirm:$false}\"",
                              pInterfaceName))) {
        SPDLOG_ERROR("Format String Error");
        return -ERR_MEMORY_STR;
    }

    if ((ret = RunCommand(cmdBuf, nullptr, 0, &retCode)) != ERR_SUCCESS) {
        SPDLOG_ERROR("Run command [{0}] error: {1}", cmdBuf, ret);
        return -ERR_CALL_SHELL;
    }

    SPDLOG_DEBUG("Run Set IP Command({1}): {0}", cmdBuf, retCode);

    return ERR_SUCCESS;
}

int SetInterfaceIpAddress(const TCHAR *pInterfaceName, const TCHAR *pIpaddr, const TCHAR *pNetmask) {
    int   ret;
    TCHAR cmdBuf[MAX_PATH];
    int   cidr;
    DWORD retCode;

    if (pInterfaceName == nullptr || lstrlen(pInterfaceName) == 0) {
        SPDLOG_ERROR("Input pInterfaceName params error: {0}", pInterfaceName);
        return -ERR_INPUT_PARAMS;
    }

    if (pIpaddr == nullptr || lstrlen(pIpaddr) == 0) {
        SPDLOG_ERROR("Input pIpaddr params error: {0}", pIpaddr);
        return -ERR_INPUT_PARAMS;
    }

    if (pNetmask == nullptr || lstrlen(pNetmask) == 0) {
        SPDLOG_WARN("Input pNetmask params error: {0}", pNetmask);
        return -ERR_INPUT_PARAMS;
    }

    cidr = NetmaskToCIDR(pNetmask);

    if (FAILED(StringCbPrintf(
            cmdBuf,
            MAX_PATH,
            "PowerShell -Command \"& {New-NetIPAddress -InterfaceAlias %s  -IPAddress %s -PrefixLength %d}\"",
            pInterfaceName,
            pIpaddr,
            cidr))) {
        SPDLOG_ERROR("Format String Error");
        return -ERR_MEMORY_STR;
    }

    if ((ret = RunCommand(cmdBuf, nullptr, 0, &retCode)) != ERR_SUCCESS) {
        SPDLOG_ERROR("Run command [{0}] error: {1}", cmdBuf, ret);
        return -ERR_CALL_SHELL;
    }

    SPDLOG_DEBUG("Run Set IP Command({1}): {0}", cmdBuf, retCode);

    return ERR_SUCCESS;
}

int SetInterfaceIpAddressFromCIDR(const TCHAR *pInterfaceName, const TCHAR *pCidrIpaddr) {
    TCHAR ipAddr[MAX_IP_LEN];
    TCHAR ip[MAX_IP_LEN];
    int   ret;
    TCHAR cmdBuf[MAX_PATH];
    DWORD retCode;

    TCHAR *token, *p = nullptr;

    if (pInterfaceName == nullptr || lstrlen(pInterfaceName) == 0) {
        SPDLOG_ERROR("Input pInterfaceName params error: {0}", pInterfaceName);
        return -ERR_INPUT_PARAMS;
    }

    if (pCidrIpaddr == nullptr || lstrlen(pCidrIpaddr) == 0) {
        SPDLOG_ERROR("Input pCidrIpaddr params error: {0}", pCidrIpaddr);
        return -ERR_INPUT_PARAMS;
    }

    StringCbCopy(ipAddr, MAX_IP_LEN, pCidrIpaddr);

    // 获取前面IP地址
    token = strtok_s(ipAddr, TEXT("/"), &p);
    if (token == nullptr) {
        SPDLOG_ERROR("CIDR IpAddress string format error: {0}", pCidrIpaddr);
        return -ERR_INPUT_PARAMS;
    }
    StringCbCopy(ip, MAX_IP_LEN, token);

    // 获取后面子网掩码
    token = strtok_s(nullptr, TEXT("/"), &p);
    if (token == nullptr) {
        SPDLOG_ERROR("CIDR IpAddress string format error: {0}", pCidrIpaddr);
        return -ERR_INPUT_PARAMS;
    }

    if (FAILED(StringCbPrintf(
            cmdBuf,
            MAX_PATH,
            "PowerShell -Command \"& {New-NetIPAddress -InterfaceAlias %s  -IPAddress %s -PrefixLength %s}\"",
            pInterfaceName,
            ip,
            token))) {
        SPDLOG_ERROR("Format String Error");
        return -ERR_MEMORY_STR;
    }

    if ((ret = RunCommand(cmdBuf, nullptr, 0, &retCode)) != ERR_SUCCESS) {
        SPDLOG_ERROR("Run command [{0}] error: {1}", cmdBuf, ret);
        return -ERR_CALL_SHELL;
    }

    SPDLOG_DEBUG("Run Set IP Command({1}): {0}", cmdBuf, retCode);

    return ERR_SUCCESS;
}

int GetAllNICInfo(PNIC_CONTENT pInfo, int *pItemCounts) {
    PIP_ADAPTER_INFO pAdapterInfo;
    DWORD            dwRetVal = 0;

    if (pItemCounts == nullptr || pInfo == nullptr) {
        SPDLOG_ERROR("Input params error: {0}");
        return -ERR_INPUT_PARAMS;
    }

    ULONG ulOutBufLen = sizeof(IP_ADAPTER_INFO);
    pAdapterInfo      = static_cast<IP_ADAPTER_INFO *>(HeapAlloc(GetProcessHeap(), 0, sizeof(IP_ADAPTER_INFO)));

    if (pAdapterInfo == nullptr) {
        SPDLOG_ERROR("Error allocating memory needed to call GetAdaptersinfo");
        return -ERR_MALLOC_MEMORY;
    }

    if (GetAdaptersInfo(pAdapterInfo, &ulOutBufLen) == ERROR_BUFFER_OVERFLOW) {
        HeapFree(GetProcessHeap(), 0, pAdapterInfo);
        pAdapterInfo = static_cast<IP_ADAPTER_INFO *>(HeapAlloc(GetProcessHeap(), 0, ulOutBufLen));
        if (pAdapterInfo == nullptr) {
            SPDLOG_ERROR("Error allocating memory needed to call GetAdaptersinfo\n");
            return -ERR_MALLOC_MEMORY;
        }
    }

    if ((dwRetVal = GetAdaptersInfo(pAdapterInfo, &ulOutBufLen)) == NO_ERROR) {
        UINT32           id       = 0;
        PIP_ADAPTER_INFO pAdapter = pAdapterInfo;
        while (pAdapter && id < NET_CARD_MAX) {
            // 网卡索引
            pInfo[id].InterfaceIndex = pAdapter->Index;
            // 网卡名称
            StringCbCopy(pInfo[id].NetCardUUID, MAX_ADAPTER_NAME_LENGTH, pAdapter->AdapterName);
            // 网卡详细描述
            StringCbCopy(pInfo[id].NetCardDescription, MAX_ADAPTER_DESCRIPTION_LENGTH, pAdapter->Description);
            // 网卡 IP 地址
            StringCbCopy(pInfo[id].NetCardIpaddr, MAX_IP_LEN - 1, pAdapter->IpAddressList.IpAddress.String);
            // 网卡子网掩码
            StringCbCopy(pInfo[id].NetCardNetmask, MAX_IP_LEN - 1, pAdapter->IpAddressList.IpMask.String);
            // 网卡网关地址
            StringCbCopy(pInfo[id].NetCardGateway, MAX_IP_LEN - 1, pAdapter->GatewayList.IpAddress.String);
            // 网卡 MAC 地址
            StringCbPrintf(pInfo[id].NetCardMacAddr,
                           20 - 1,
                           "%02X:%02X:%02X:%02X:%02X:%02X",
                           pAdapter->Address[0],
                           pAdapter->Address[1],
                           pAdapter->Address[2],
                           pAdapter->Address[3],
                           pAdapter->Address[4],
                           pAdapter->Address[5]);

            id++;
            pAdapter = pAdapter->Next;
        }

        *pItemCounts = id;
    } else {
        SPDLOG_ERROR("GetAdaptersInfo failed with error: {0}\n", dwRetVal);
        return -ERR_SYS_CALL;
    }

    if (pAdapterInfo) {
        HeapFree(GetProcessHeap(), 0, pAdapterInfo);
    }

    return ERR_SUCCESS;
}

#if 0
int GetAllNICInfo(PNIC_CONTENT pInfo, int *pItemCounts) {
    int              nRel, id = 0;
    IP_ADDR_STRING  *pIpAddrString;
    PIP_ADAPTER_INFO pIpAdapterInfo;
    unsigned long    stSize;

    if (pItemCounts == nullptr || pInfo == nullptr) {
        return -ERR_INPUT_PARAMS;
    }

    pIpAdapterInfo = new IP_ADAPTER_INFO[NET_CARD_MAX];
    stSize         = sizeof(IP_ADAPTER_INFO) * NET_CARD_MAX;

    // WIN32 API get net card information
    nRel = GetAdaptersInfo(pIpAdapterInfo, &stSize);

    if (ERROR_BUFFER_OVERFLOW == nRel) {
        delete[] pIpAdapterInfo;
        return -ERR_MALLOC_MEMORY;
    }

    PIP_ADAPTER_INFO cur = pIpAdapterInfo;

    while (cur) {
        pInfo[id].InterfaceIndex = cur->Index;
        StringCbCopy(pInfo[id].NetCardUUID, MAX_ADAPTER_NAME_LENGTH, cur->AdapterName);
        StringCbCopy(pInfo[id].NetCardDescription, MAX_ADAPTER_DESCRIPTION_LENGTH, cur->Description);

        switch (cur->Type) {
            case MIB_IF_TYPE_ETHERNET:
                pIpAddrString = &(cur->IpAddressList);
                StringCbCopy(pInfo[id].NetCardIpaddr, MAX_IP_LEN - 1, pIpAddrString->IpAddress.String);
                StringCbCopy(pInfo[id].NetCardNetmask, MAX_IP_LEN - 1, pIpAddrString->IpMask.String);
                break;
            case MIB_IF_TYPE_OTHER:
            case MIB_IF_TYPE_TOKENRING:
            case MIB_IF_TYPE_FDDI:
            case MIB_IF_TYPE_PPP:
            case MIB_IF_TYPE_LOOPBACK:
            case MIB_IF_TYPE_SLIP:
                break;
            default:    // WIFI ,Unknown type
                pIpAddrString = &(cur->IpAddressList);
                StringCbCopy(pInfo[id].NetCardIpaddr, MAX_IP_LEN - 1, pIpAddrString->IpAddress.String);
                StringCbCopy(pInfo[id].NetCardNetmask, MAX_IP_LEN - 1, pIpAddrString->IpMask.String);
                break;
        }

        StringCbPrintf(pInfo[id].NetCardMacAddr,
                       20 - 1,
                       "%02X:%02X:%02X:%02X:%02X:%02X",
                       cur->Address[0],
                       cur->Address[1],
                       cur->Address[2],
                       cur->Address[3],
                       cur->Address[4],
                       cur->Address[5]);

        id++;
        cur = cur->Next;
    }

    *pItemCounts = id;

    delete[] pIpAdapterInfo;

    return ERR_SUCCESS;
}
#endif