#include "pch.h"
#include <strsafe.h>
#include <winsock2.h>
#include <ws2ipdef.h>
#include <ws2tcpip.h>
#include <spdlog/spdlog.h>

#include "usrerr.h"
#include "misc.h"
#include "tunnel.h"

static const TCHAR *p2_table(unsigned pow) {
    static const TCHAR *pow2[] = {
        TEXT("1"),
        TEXT("2"),
        TEXT("4"),
        TEXT("8"),
        TEXT("16"),
        TEXT("32"),
        TEXT("64"),
        TEXT("128"),
        TEXT("256"),
        TEXT("512"),
        TEXT("1024"),
        TEXT("2048"),
        TEXT("4096"),
        TEXT("8192"),
        TEXT("16384"),
        TEXT("32768"),
        TEXT("65536"),
        TEXT("131072"),
        TEXT("262144"),
        TEXT("524288"),
        TEXT("1048576"),
        TEXT("2097152"),
        TEXT("4194304"),
        TEXT("8388608"),
        TEXT("16777216"),
        TEXT("33554432"),
        TEXT("67108864"),
        TEXT("134217728"),
        TEXT("268435456"),
        TEXT("536870912"),
        TEXT("1073741824"),
        TEXT("2147483648"),
        TEXT("4294967296"),
        TEXT("8589934592"),
        TEXT("17179869184"),
        TEXT("34359738368"),
        TEXT("68719476736"),
        TEXT("137438953472"),
        TEXT("274877906944"),
        TEXT("549755813888"),
        TEXT("1099511627776"),
        TEXT("2199023255552"),
        TEXT("4398046511104"),
        TEXT("8796093022208"),
        TEXT("17592186044416"),
        TEXT("35184372088832"),
        TEXT("70368744177664"),
        TEXT("140737488355328"),
        TEXT("281474976710656"),
        TEXT("562949953421312"),
        TEXT("1125899906842624"),
        TEXT("2251799813685248"),
        TEXT("4503599627370496"),
        TEXT("9007199254740992"),
        TEXT("18014398509481984"),
        TEXT("36028797018963968"),
        TEXT("72057594037927936"),
        TEXT("144115188075855872"),
        TEXT("288230376151711744"),
        TEXT("576460752303423488"),
        TEXT("1152921504606846976"),
        TEXT("2305843009213693952"),
        TEXT("4611686018427387904"),
        TEXT("9223372036854775808"),
        TEXT("18446744073709551616"),
        TEXT("36893488147419103232"),
        TEXT("73786976294838206464"),
        TEXT("147573952589676412928"),
        TEXT("295147905179352825856"),
        TEXT("590295810358705651712"),
        TEXT("1180591620717411303424"),
        TEXT("2361183241434822606848"),
        TEXT("4722366482869645213696"),
        TEXT("9444732965739290427392"),
        TEXT("18889465931478580854784"),
        TEXT("37778931862957161709568"),
        TEXT("75557863725914323419136"),
        TEXT("151115727451828646838272"),
        TEXT("302231454903657293676544"),
        TEXT("604462909807314587353088"),
        TEXT("1208925819614629174706176"),
        TEXT("2417851639229258349412352"),
        TEXT("4835703278458516698824704"),
        TEXT("9671406556917033397649408"),
        TEXT("19342813113834066795298816"),
        TEXT("38685626227668133590597632"),
        TEXT("77371252455336267181195264"),
        TEXT("154742504910672534362390528"),
        TEXT("309485009821345068724781056"),
        TEXT("618970019642690137449562112"),
        TEXT("1237940039285380274899124224"),
        TEXT("2475880078570760549798248448"),
        TEXT("4951760157141521099596496896"),
        TEXT("9903520314283042199192993792"),
        TEXT("19807040628566084398385987584"),
        TEXT("39614081257132168796771975168"),
        TEXT("79228162514264337593543950336"),
        TEXT("158456325028528675187087900672"),
        TEXT("316912650057057350374175801344"),
        TEXT("633825300114114700748351602688"),
        TEXT("1267650600228229401496703205376"),
        TEXT("2535301200456458802993406410752"),
        TEXT("5070602400912917605986812821504"),
        TEXT("10141204801825835211973625643008"),
        TEXT("20282409603651670423947251286016"),
        TEXT("40564819207303340847894502572032"),
        TEXT("81129638414606681695789005144064"),
        TEXT("162259276829213363391578010288128"),
        TEXT("324518553658426726783156020576256"),
        TEXT("649037107316853453566312041152512"),
        TEXT("1298074214633706907132624082305024"),
        TEXT("2596148429267413814265248164610048"),
        TEXT("5192296858534827628530496329220096"),
        TEXT("10384593717069655257060992658440192"),
        TEXT("20769187434139310514121985316880384"),
        TEXT("41538374868278621028243970633760768"),
        TEXT("83076749736557242056487941267521536"),
        TEXT("166153499473114484112975882535043072"),
        TEXT("332306998946228968225951765070086144"),
        TEXT("664613997892457936451903530140172288"),
        TEXT("1329227995784915872903807060280344576"),
        TEXT("2658455991569831745807614120560689152"),
        TEXT("5316911983139663491615228241121378304"),
        TEXT("10633823966279326983230456482242756608"),
        TEXT("21267647932558653966460912964485513216"),
        TEXT("42535295865117307932921825928971026432"),
        TEXT("85070591730234615865843651857942052864"),
        TEXT("170141183460469231731687303715884105728"),
    };
    if (pow <= 127) {
        return pow2[pow];
    }
    return TEXT("");
}

static int vasprintf(TCHAR **strp, const TCHAR *fmt, va_list ap) {
    // _vscprintf tells you how big the buffer needs to be
    const int len = _vscprintf(fmt, ap);
    if (len == -1) {
        return -1;
    }
    const size_t size = static_cast<size_t>(len) + 1;
    const auto   str  = static_cast<TCHAR *>(malloc(size));
    if (!str) {
        return -1;
    }
    // _vsprintf_s is the "secure" version of vsprintf
    const int r = vsprintf_s(str, len + 1, fmt, ap);
    if (r == -1) {
        free(str);
        return -1;
    }
    *strp = str;
    return r;
}

static int asprintf(TCHAR **strp, const TCHAR *fmt, ...) {
    va_list ap;
    va_start(ap, fmt);
    const int r = vasprintf(strp, fmt, ap);
    va_end(ap);
    return r;
}

static int bit_count(unsigned int i) {
    int          c        = 0;
    unsigned int seen_one = 0;

    while (i > 0) {
        if (i & 1) {
            seen_one = 1;
            c++;
        } else {
            if (seen_one) {
                return -1;
            }
        }
        i >>= 1;
    }

    return c;
}

/**
 * @brief creates a netmask from a specified number of bits
 * This function converts a prefix length to a netmask.  As CIDR (classless
 * internet domain internet domain routing) has taken off, more an more IP
 * addresses are being specified in the format address/prefix
 * (i.e. 192.168.2.3/24, with a corresponding netmask 255.255.255.0).  If you
 * need to see what netmask corresponds to the prefix part of the address, this
 * is the function.  See also @ref mask2prefix.
 * @param prefix prefix is the number of bits to create a mask for.
 * @return a network mask, in network byte order.
 */
unsigned int prefix2mask(int prefix) {
    if (prefix) {
        return htonl(~((1 << (32 - prefix)) - 1));
    } else {
        return htonl(0);
    }
}

/**
* @brief calculates the number of bits masked off by a netmask.
* This function calculates the significant bits in an IP address as specified by
* a netmask.  See also @ref prefix2mask.
* @param mask is the netmask, specified as an struct in_addr in network byte order.
* @return the number of significant bits.
*/
int mask2prefix(IN_ADDR mask) {
    return bit_count(ntohl(mask.s_addr));
}

static int ipv4_mask_to_int(const char *prefix) {
    int     ret;
    IN_ADDR in;

    ret = inet_pton(AF_INET, prefix, &in);
    if (ret == 0) {
        return -1;
    }

    return mask2prefix(in);
}

/**
*  @brief calculate broadcast address given an IP address and a prefix length.
*  @param addr an IP address in network byte order.
*  @param prefix a prefix length.
*  @return the calculated broadcast address for the network, in network byte order.
*/
static IN_ADDR calc_broadcast(IN_ADDR addr, int prefix) {
    IN_ADDR mask;
    IN_ADDR broadcast;

    mask.s_addr = prefix2mask(prefix);

    memset(&broadcast, 0, sizeof(broadcast));
    broadcast.s_addr = (addr.s_addr & mask.s_addr) | ~mask.s_addr;
    return broadcast;
}

/**
* @brief calculates the network address for a specified address and prefix.
* @param addr an IP address, in network byte order
* @param prefix the network prefix
* @return the base address of the network that addr is associated with, in
*  network byte order.
*/
static IN_ADDR calc_network(IN_ADDR addr, int prefix) {
    IN_ADDR mask;
    IN_ADDR network;

    mask.s_addr = prefix2mask(prefix);

    memset(&network, 0, sizeof(network));
    network.s_addr = addr.s_addr & mask.s_addr;
    return network;
}

static TCHAR *ipv4_prefix_to_hosts(TCHAR *hosts, unsigned hosts_size, unsigned prefix) {
    if (prefix >= 31) {
        StringCbPrintf(hosts, hosts_size, TEXT("%s"), p2_table(32 - prefix));
    } else {
        unsigned int tmp;
        tmp = (1 << (32 - prefix)) - 2;
        StringCbPrintf(hosts, hosts_size, TEXT("%u"), tmp);
    }
    return hosts;
}

static int str_to_prefix(int *ipv6, const char *prefixStr, unsigned fix) {
    int prefix;

    if (!(*ipv6) && strchr(prefixStr, '.')) { /* prefix is 255.x.x.x */
        prefix = ipv4_mask_to_int(prefixStr);
    } else {
        prefix = strtol(prefixStr, nullptr, 10);
    }

    if (fix && (prefix > 32 && !(*ipv6))) {
        *ipv6 = 1;
    }

    if (prefix < 0 || (((*ipv6) && prefix > 128) || (!(*ipv6) && prefix > 32))) {
        return -1;
    }
    return prefix;
}

static int GetIpV4Info(const TCHAR *pIpStr, int prefix, PIP_INFO pInfo, unsigned int flags) {
    IN_ADDR ip, netmask, network, broadcast, minhost, maxhost;
    TCHAR   namebuf[INET_ADDRSTRLEN + 1];
    TCHAR  *ipStr = _strdup(pIpStr);

    memset(pInfo, 0, sizeof(*pInfo));

    if (inet_pton(AF_INET, ipStr, &ip) <= 0) {
        SPDLOG_ERROR(TEXT("ipcalc: bad IPv4 address: {0}"), ipStr);
        free(ipStr);
        return -ERR_UN_SUPPORT;
    }

    /* Handle CIDR entries such as 172/8 */
    if (prefix >= 0) {
        auto tmp = const_cast<TCHAR *>(ipStr);
        int  i;

        for (i = 3; i > 0; i--) {
            tmp = strchr(tmp, '.');
            if (!tmp) {
                break;
            } else {
                tmp++;
            }
        }

        tmp = nullptr;
        for (; i > 0; i--) {
            if (asprintf(&tmp, "%s.0", ipStr) == -1) {
                SPDLOG_ERROR(TEXT("Memory allocation failure"));
                free(ipStr);
                return -ERR_MALLOC_MEMORY;
            }
            ipStr = tmp;
        }
    } else {    // assume good old days classful Internet
        prefix = 32;
    }

    if (prefix > 32) {
        SPDLOG_ERROR(TEXT("ipcalc: bad IPv4 prefix: {0}"), prefix);
        free(ipStr);
        return -ERR_UN_SUPPORT;
    }

    if (inet_ntop(AF_INET, &ip, namebuf, sizeof(namebuf)) == 0) {
        SPDLOG_ERROR(TEXT("ipcalc: error calculating the IPv4 network"));
        free(ipStr);
        return -ERR_UN_SUPPORT;
    }
    StringCbCopy(pInfo->ip, MAX_IP_LEN, namebuf);

    netmask.s_addr = prefix2mask(prefix);
    memset(namebuf, '\0', sizeof(namebuf));

    if (inet_ntop(AF_INET, &netmask, namebuf, INET_ADDRSTRLEN) == nullptr) {
        SPDLOG_ERROR(TEXT("inet_ntop error"));
        free(ipStr);
        return -ERR_UN_SUPPORT;
    }
    StringCbCopy(pInfo->netmask, MAX_IP_LEN, namebuf);

    pInfo->prefix = prefix;

    broadcast = calc_broadcast(ip, prefix);

    memset(namebuf, '\0', sizeof(namebuf));
    if (inet_ntop(AF_INET, &broadcast, namebuf, INET_ADDRSTRLEN) == nullptr) {
        SPDLOG_ERROR(TEXT("inet_ntop error"));
        free(ipStr);
        return -ERR_UN_SUPPORT;
    }

    StringCbCopy(pInfo->broadcast, MAX_IP_LEN, namebuf);

    network = calc_network(ip, prefix);

    memset(namebuf, '\0', sizeof(namebuf));
    if (inet_ntop(AF_INET, &network, namebuf, INET_ADDRSTRLEN) == nullptr) {
        SPDLOG_ERROR(TEXT("inet_ntop error"));
        free(ipStr);
        return -ERR_UN_SUPPORT;
    }

    StringCbCopy(pInfo->network, MAX_IP_LEN, namebuf);

    if (prefix < 32) {
        memcpy(&minhost, &network, sizeof(minhost));

        if (prefix <= 30) {
            minhost.s_addr = htonl(ntohl(minhost.s_addr) | 1);
        }
        if (inet_ntop(AF_INET, &minhost, namebuf, INET_ADDRSTRLEN) == nullptr) {
            SPDLOG_ERROR(TEXT("inet_ntop error"));
            free(ipStr);
            return -ERR_UN_SUPPORT;
        }
        StringCbCopy(pInfo->hostmin, MAX_IP_LEN, namebuf);

        memcpy(&maxhost, &network, sizeof(minhost));
        maxhost.s_addr |= ~netmask.s_addr;
        if (prefix <= 30) {
            maxhost.s_addr = htonl(ntohl(maxhost.s_addr) - 1);
        }
        if (inet_ntop(AF_INET, &maxhost, namebuf, sizeof(namebuf)) == 0) {
            SPDLOG_ERROR(TEXT("ipcalc: error calculating the IPv4 network"));
            free(ipStr);
            return -ERR_UN_SUPPORT;
        }

        StringCbCopy(pInfo->hostmax, MAX_IP_LEN, namebuf);
    } else {
        StringCbCopy(pInfo->hostmin, MAX_IP_LEN, pInfo->network);
        StringCbCopy(pInfo->hostmax, MAX_IP_LEN, pInfo->network);
    }

    ipv4_prefix_to_hosts(pInfo->hosts, sizeof(pInfo->hosts), prefix);

    free(ipStr);
    return ERR_SUCCESS;
}

/**
 * @brief 计算 IPv4 网络信息
 * @param[in] pIpStr IPv4 网络信息 '/' 分割，支持CIDR以及子网掩码 example: 192.168.1.32/24, 192.168.1.32/255.255.255.0
 * @param[out] pInfo 计算结果
 * @return 0: 成功， 小于0 失败 @see USER_ERRNO
 * - -ERR_INPUT_PARAMS 输入参数错误
 * - -ERR_UN_SUPPORT 不支持的格式转换
 * - -ERR_MALLOC_MEMORY 分配内存失败
 * - ERR_SUCCESS 成功
 */
int GetIpV4InfoFromCIDR(const TCHAR *pIpStr, PIP_INFO pInfo) {
    int    ret, prefix, familyIPv6 = 0;
    TCHAR *prefixStr;
    TCHAR *ipStr = _strdup(pIpStr);

    if (pIpStr == nullptr || lstrlen(pIpStr) < MIN_IP_LEN || lstrlen(pIpStr) >= MAX_IP_LEN) {
        SPDLOG_ERROR(TEXT("Input pIpStr format error: {}."), pIpStr);
        return -ERR_INPUT_PARAMS;
    }

    if (pInfo == nullptr) {
        SPDLOG_ERROR(TEXT("Input pInfo is NULL."));
        return -ERR_INPUT_PARAMS;
    }

    if (strchr(ipStr, '/') != nullptr) {
        prefixStr  = static_cast<TCHAR *>(strchr(ipStr, '/'));
        *prefixStr = '\0'; /* fix up ipStr */
        prefixStr++;
    } else {
        SPDLOG_ERROR(TEXT("Input pIpStr isn't CIDR format: {}."), pIpStr);
        free(ipStr);
        return -ERR_INPUT_PARAMS;
    }

    if (strchr(prefixStr, '.') != nullptr) {
        prefix = ipv4_mask_to_int(prefixStr);
    } else {
        prefix = str_to_prefix(&familyIPv6, prefixStr, 0);
    }

    ret = GetIpV4Info(ipStr, prefix, pInfo, 0);

    free(ipStr);
    return ret;
}

/**
 * @brief 计算 IPv4 网络信息
 * @param[in] pIpStr IPv4 地址
 * @param[in] pNetmask IPv4子网掩码
 * @param[out] pInfo 计算结果
 * @return 成功， 小于0 失败 @see USER_ERRNO
 * - -ERR_INPUT_PARAMS 输入参数错误
 * - -ERR_UN_SUPPORT 不支持的格式转换
 * - -ERR_MALLOC_MEMORY 分配内存失败
 * - ERR_SUCCESS 成功
 */
int GetIpV4InfoFromNetmask(const TCHAR *pIpStr, const TCHAR *pNetmask, PIP_INFO pInfo) {
    int prefix;

    if (pIpStr == nullptr || lstrlen(pIpStr) < MIN_IP_LEN || lstrlen(pIpStr) >= MAX_IP_LEN) {
        SPDLOG_ERROR(TEXT("Input pIpStr format error: {}."), pIpStr);
        return -ERR_INPUT_PARAMS;
    }

    if (pNetmask == nullptr || lstrlen(pNetmask) < MIN_IP_LEN || lstrlen(pNetmask) >= MAX_IP_LEN) {
        SPDLOG_ERROR(TEXT("Input pNetmask format error: {}."), pNetmask);
        return -ERR_INPUT_PARAMS;
    }

    if (pInfo == nullptr) {
        SPDLOG_ERROR(TEXT("Input pInfo is NULL."));
        return -ERR_INPUT_PARAMS;
    }

    prefix = ipv4_mask_to_int(pNetmask);
    return GetIpV4Info(pIpStr, prefix, pInfo, 0);
}

int GetIpV4InfoFromHostname(int family, const char *host, PIP_INFO pInfo) {
    addrinfo   *res, *rp;
    addrinfo    hints {};
    int         err;
    static char ipname[64];
    void       *addr;

    memset(&hints, 0, sizeof(hints));
    hints.ai_family = family;

    err = getaddrinfo(host, nullptr, &hints, &res);
    if (err != 0) {
        return -ERR_INPUT_PARAMS;
    }

    for (rp = res; rp != nullptr; rp = rp->ai_next) {
        if (rp->ai_family == AF_INET) {
            addr = (&reinterpret_cast<sockaddr_in *>(rp->ai_addr)->sin_addr);
        } else {
            addr = (&reinterpret_cast<sockaddr_in6 *>(rp->ai_addr)->sin6_addr);
        }

        if (inet_ntop(rp->ai_family, addr, ipname, sizeof(ipname)) != nullptr) {
            freeaddrinfo(res);
            StringCbCopy(pInfo->hostip, MAX_IP_LEN, ipname);
            return ERR_SUCCESS;
        }
    }

    freeaddrinfo(res);
    return ERR_ITEM_EXISTS;
}