//
// Created by xajhuang on 2023/4/20.
//
#include <uv.h>
#include <string.h>
#include <sys/time.h>
#include "user_errno.h"
#include "dhcp_network.h"
#include "dhcp_options.h"
#include "zlog_module.h"
#include "main.h"
#include "rfc2131.h"
#include "misc.h"

static U8    g_dhcpReqParams[] = {0x01, 0x1c, 0x02, 0x03, 0x0f, 0x06, 0x77, 0x0c, 0x2c, 0x2f, 0x1a, 0x79, 0x2a};
static char *g_pNicName        = NULL;

static PDHCP_PROTO fill_dhcp_pkg(U8 *pRsp, PDHCP_INFO pInfo) {
    struct ethhdr *pEth  = (struct ethhdr *)pRsp;
    struct iphdr  *pIp   = NULL;
    struct udphdr *pUdp  = NULL;
    PDHCP_PROTO    pDhcp = NULL;

    memset(pRsp, 0, MAX_DHCP_PKG_SIZE);

    // 目的地 MAC 地址
    memset(pEth->h_dest, 0xFF, ETH_ALEN);
    // 源 MAC 地址
    memcpy(pEth->h_source, pInfo->mac, ETH_ALEN);

    if (pInfo->vni == 0) {
        pIp           = (struct iphdr *)(pRsp + IP_HDR_OFFSET);
        pUdp          = (struct udphdr *)(UDP_HDR_OFFSET);
        pDhcp         = (PDHCP_PROTO)((U8 *)pRsp + DHCP_OFFSET);
        pEth->h_proto = htons(ETH_P_IP);
    } else if (pInfo->vni >= 1 && pInfo->vni <= 4094) {
        PVLAN_HDR pHdr = (PVLAN_HDR)(pRsp + sizeof(struct ethhdr));
        pIp            = (struct iphdr *)(pRsp + IP_HDR_OFFSET + sizeof(VLAN_HDR));
        pUdp           = (struct udphdr *)(pRsp + UDP_HDR_OFFSET + sizeof(VLAN_HDR));
        pDhcp          = (PDHCP_PROTO)((U8 *)pRsp + DHCP_OFFSET + sizeof(VLAN_HDR));
        pEth->h_proto  = htons(ETH_P_8021Q);
        pHdr->id       = htons(pInfo->vni);
        pHdr->type     = htons(ETH_P_IP);
    } else {
        PVLAN_HDR2 pHdr = (PVLAN_HDR2)(pRsp + sizeof(struct ethhdr));
        pIp             = (struct iphdr *)(pRsp + IP_HDR_OFFSET + sizeof(VLAN_HDR2));
        pUdp            = (struct udphdr *)(pRsp + UDP_HDR_OFFSET + sizeof(VLAN_HDR2));
        pDhcp           = (PDHCP_PROTO)((U8 *)pRsp + DHCP_OFFSET + sizeof(VLAN_HDR2));
        pEth->h_proto   = htons(ETH_P_8021Q);
        pHdr->id1       = htons(4094);
        pHdr->h_type    = htons(ETH_P_8021Q);
        pHdr->id2       = htons((pInfo->vni - 4094) & 0x0FFF);
        pHdr->type      = htons(ETH_P_IP);
    }

    // IP 头
    pIp->version  = IPVERSION;
    pIp->ihl      = 5;
    pIp->tos      = 0;
    pIp->tot_len  = 0;
    pIp->id       = 0;
    pIp->frag_off = 0;
    pIp->ttl      = 128;
    pIp->protocol = IPPROTO_UDP;
    pIp->check    = 0;
    pIp->saddr    = INADDR_ANY;
    pIp->daddr    = INADDR_BROADCAST;

    // UDP 头
    pUdp->source = htons(DHCP_CLI_PORT);
    pUdp->dest   = htons(DHCP_SVR_PORT);
    pUdp->len    = 0;
    pUdp->check  = 0;

    // DHCP 协议内容
    pDhcp->op     = BOOTP_REQUEST;
    pDhcp->htype  = 0x01;
    pDhcp->hlen   = ETH_ALEN;
    pDhcp->hops   = 0;
    pDhcp->xid    = 0;
    pDhcp->secs   = 0;
    pDhcp->flags  = 0;
    pDhcp->ciaddr = INADDR_ANY;
    pDhcp->yiaddr = INADDR_ANY;
    pDhcp->siaddr = INADDR_ANY;
    pDhcp->giaddr = INADDR_ANY;
    memcpy(pDhcp->chaddr, pInfo->mac, ETH_ALEN);
    pDhcp->cookie = htonl(DHCP_COOKIE_VAL);

    return pDhcp;
}

static U32 dhcp_pkk_checksum(U8 *pRsp, U32 optSize) {
    struct iphdr  *pIp  = NULL;
    struct udphdr *pUdp = NULL;
    U16            csum;
    U32            tolSize;
    U8             vlanSize = 0;

    switch (get_package_vlan_type(pRsp)) {
        default:
            pIp     = (struct iphdr *)(pRsp + IP_HDR_OFFSET);
            pUdp    = (struct udphdr *)(UDP_HDR_OFFSET);
            // 计算包总长度
            tolSize = optSize + sizeof(DHCP_PROTO) + DHCP_OFFSET;
            break;
        case VLAN_LEVEL1:
            pIp      = (struct iphdr *)(pRsp + IP_HDR_OFFSET + sizeof(VLAN_HDR));
            pUdp     = (struct udphdr *)(pRsp + UDP_HDR_OFFSET + sizeof(VLAN_HDR));
            // 计算包总长度
            tolSize  = optSize + sizeof(DHCP_PROTO) + DHCP_OFFSET + sizeof(VLAN_HDR);
            vlanSize = sizeof(VLAN_HDR);
            break;
        case VLAN_LEVEL2:
            pIp      = (struct iphdr *)(pRsp + IP_HDR_OFFSET + sizeof(VLAN_HDR2));
            pUdp     = (struct udphdr *)(pRsp + UDP_HDR_OFFSET + sizeof(VLAN_HDR2));
            // 计算包总长度
            tolSize  = optSize + sizeof(DHCP_PROTO) + DHCP_OFFSET + sizeof(VLAN_HDR2);
            vlanSize = sizeof(VLAN_HDR2);
            break;
    }

    // 计算 IP 数据长度
    pIp->tot_len = htons(tolSize - sizeof(struct ethhdr) - vlanSize);
    // 计算 UDP 数据长度
    pUdp->len    = htons(tolSize - sizeof(struct ethhdr) - vlanSize - sizeof(struct iphdr));

    // 计算 IP 校验和
    csum       = htons(ip_checksum((unsigned char *)pIp));
    pIp->check = htons(csum);

    // 计算 UDP 校验和
    csum        = htons(udp_checksum(pIp->saddr, pIp->daddr, (unsigned char *)pUdp));
    pUdp->check = htons(csum);

    return tolSize;
}

U8 *dhcp_create_request_req(PDHCP_INFO pInfo, int *pOutSize) {
    U8 *pOpt;
    U8 *pReqData = (U8 *)malloc(MAX_DHCP_PKG_SIZE);

    if (pReqData) {
        //PDHCP_PACKAGE pDhcp = (PDHCP_PACKAGE)pReqData;
        PDHCP_PROTO pDhcp = fill_dhcp_pkg(pReqData, pInfo);

        //pkg_init_head(pDhcp, pInfo);

        pDhcp->xid = htonl(pInfo->offerRsp.xid);

        // DHCP Options
        pOpt = pDhcp->options;

        // DHCP 消息类型
        pOpt  += dhcp_add_u8_option(pOpt, OPT_MESSAGETYPE, DHCP_MSG_REQUEST);
        // DHCP 请求参数列表
        pOpt  += dhcp_add_buf_option(pOpt, OPT_PARAMREQLIST, g_dhcpReqParams, 13);
        // DHCP 主机名
        pOpt  += dhcp_add_string_option(pOpt, OPT_HOSTNAME, pInfo->hostname);
        // DHCP Server 主机
        pOpt  += dhcp_add_u32_option(pOpt, OPT_SERVERID, htonl(pInfo->offerRsp.svrIp));
        // DHCP 请求分配IP
        pOpt  += dhcp_add_u32_option(pOpt, OPT_REQUESTEDIPADDR, htonl(pInfo->offerRsp.ipAddr));
        // 结束
        *pOpt  = OPT_END;

        *pOutSize = (int)dhcp_pkk_checksum(pReqData, (int)((pOpt - pDhcp->options) + 1));
    }

    return pReqData;
}

U8 *dhcp_create_discover_req(PDHCP_INFO pInfo, int *pOutSize) {
    U8 *pOpt;
    int tolSize;
    U8 *pReqData = (U8 *)malloc(MAX_DHCP_PKG_SIZE);

    if (pReqData) {
        //PDHCP_PACKAGE pDhcp = (PDHCP_PACKAGE)pReqData;
        PDHCP_PROTO pDhcp = fill_dhcp_pkg(pReqData, pInfo);
        //memset(pReqData, 0, MAX_DHCP_PKG_SIZE);

        //pkg_init_head(pDhcp, pInfo);
        pDhcp->xid = htonl((rand_number() & 0xFF000000) + pInfo->index);

        // DHCP Options
        pOpt = pDhcp->options;

        // DHCP 消息类型
        pOpt  += dhcp_add_u8_option(pOpt, OPT_MESSAGETYPE, DHCP_MSG_DISCOVER);
        // DHCP 主机名
        pOpt  += dhcp_add_string_option(pOpt, OPT_HOSTNAME, pInfo->hostname);
        // DHCP 请求参数列表
        pOpt  += dhcp_add_buf_option(pOpt, OPT_PARAMREQLIST, g_dhcpReqParams, 13);
        // 结束
        *pOpt  = OPT_END;

        *pOutSize = (int)dhcp_pkk_checksum(pReqData, (int)((pOpt - pDhcp->options) + 1));
    }

    return pReqData;
}

static void on_dhcp_recv(uv_work_t *req) {
    int               status = 0;
    PDHCP_INFO        pInfo;
    DHCP_RSP          rspDhcp;
    DHCP_OPT          optMsg, opt;
    int               ret;
    PPKG_PROCESS_INFO pWork   = (PPKG_PROCESS_INFO)req->data;
    U32               optSize = 0;
    PDHCP_PROTO       pDhcp   = get_dhcp_date(pWork->pPkgBase, &optSize, pWork->nSize);
    U32               id      = DHCP_XID(pDhcp->xid);

    // Check op flag
    if (pDhcp->op != BOOTP_REPLY) {
        LOG_MOD(error, ZM_DHCP_NET, "Error message op code %d\n", pDhcp->op);
        return;
    }

    // 获取消息类型
    ret = dhcp_get_option(OPT_MESSAGETYPE, pDhcp->options, optSize, &optMsg);
    if (ret != ERR_SUCCESS) {
        LOG_MOD(error, ZM_DHCP_NET, "Get \'message type\' option error %d\n", ret);
        return;
    }

    pInfo = get_dhcp_info_by_id(id);

    if (pInfo == NULL) {
        LOG_MOD(error, ZM_DHCP_NET, "Unknown Client %d\n", DHCP_XID(pDhcp->xid) & 0xFFFFFF);
        return;
    }

    memset(&rspDhcp, 0, sizeof(DHCP_RSP));
    rspDhcp.xid    = id;
    rspDhcp.ipAddr = ntohl(pDhcp->yiaddr);
    memcpy(rspDhcp.cliMac, pDhcp->chaddr, ETH_ALEN);
    memcpy(rspDhcp.svrHostname, pDhcp->sname, 64);

    switch (*optMsg.pValue) {
        case DHCP_MSG_OFFER:
            ret = dhcp_get_option(OPT_NETMASK, pDhcp->options, optSize, &opt);
            if (ret == ERR_SUCCESS && opt.len == sizeof(U32)) {
                rspDhcp.netmask = ntohl(*((U32 *)opt.pValue));
            }

            ret = dhcp_get_option(OPT_DNS, pDhcp->options, optSize, &opt);
            if (ret == ERR_SUCCESS) {
                U32 *pVal        = (U32 *)opt.pValue;
                rspDhcp.primeDNS = ntohl(*pVal);
                if (optSize == 8) {
                    pVal++;
                    rspDhcp.secondDNS = ntohl(*pVal);
                }
            }

            ret = dhcp_get_option(OPT_ROUTER, pDhcp->options, optSize, &opt);
            if (ret == ERR_SUCCESS) {
                rspDhcp.route = ntohl(*((U32 *)opt.pValue));
            }

            ret = dhcp_get_option(OPT_IPADDRLEASE, pDhcp->options, optSize, &opt);
            if (ret == ERR_SUCCESS) {
                rspDhcp.leaseTime = ntohl(*((U32 *)opt.pValue));
            }

            ret = dhcp_get_option(OPT_SERVERID, pDhcp->options, optSize, &opt);
            if (ret == ERR_SUCCESS) {
                rspDhcp.svrIp = ntohl(*((U32 *)opt.pValue));
            }

            ret = dhcp_get_option(OPT_DOMAINNAME, pDhcp->options, optSize, &opt);
            if (ret == ERR_SUCCESS) {
                strncpy(rspDhcp.domainName, (char *)opt.pValue, MIN((int)opt.len, 64));
            }

            cacheDhcpOfferBuffer(pInfo, pWork->pPkgBase, pWork->nSize);
            memcpy(&pInfo->offerRsp, &rspDhcp, sizeof(DHCP_RSP));
            gettimeofday(&pInfo->pOfferBuf.tm, NULL);

            pInfo->step   = STEP_OFFER;
            pInfo->status = STA_RECV_RSP;
            break;
        case DHCP_MSG_ACK:
            cacheDhcpAckBuffer(pInfo, pWork->pPkgBase, pWork->nSize);
            memcpy(&pInfo->ackRsp, &rspDhcp, sizeof(DHCP_RSP));
            gettimeofday(&pInfo->pAckBuf.tm, NULL);
            pInfo->step   = STEP_ACK;
            pInfo->status = STA_RECV_RSP;
            break;
        case DHCP_MSG_NAK:
            cacheDhcpAckBuffer(pInfo, pWork->pPkgBase, pWork->nSize);
            memcpy(&pInfo->ackRsp, &rspDhcp, sizeof(DHCP_RSP));
            gettimeofday(&pInfo->pAckBuf.tm, NULL);
            pInfo->status = STA_RECV_NACK;
            break;
    }
}

int dhcp_tools_init_network(const char *pNicName) {
    int ret;

    if (g_pNicName == NULL) {
        g_pNicName = strdup(pNicName);
    } else if (strcmp(pNicName, g_pNicName) != 0) {
        dhcp_uninit();
        free(g_pNicName);
        g_pNicName = strdup(pNicName);
    } else {
        return ERR_SUCCESS;
    }

    init_filter("vlan or (udp and dst port 68)");

    ret = create_udp_raw_socket(g_pNicName);
    if (ret != ERR_SUCCESS) {
        LOG_MOD(error, ZM_DHCP_NET, "Create receive RAW Socket Error: %s(%d)\n", getErrorEnumNameString(-ret), ret);
        return ret;
    }

    dhcp_option_cfg_init();
    init_raw_socket_poll(on_dhcp_recv, NULL);

    return ERR_SUCCESS;
}