/*
Copyright (C) 2020 Simon Wunderlich, Marek Sobe
Copyright (C) 2020 Doodle Labs

SPDX-License-Identifier: Apache-2.0

Open Drone ID C Library

Maintainer:
Simon Wunderlich
sw@simonwunderlich.de
*/
/*该文件是对开源wifi修改适配国产标准*/
#if defined(ARDUINO_ARCH_ESP32)
#include <Arduino.h>
int clock_gettime(clockid_t, struct timespec *);
#else 
#include <string.h>
#include <stddef.h>
#include <stdio.h>
#endif

#include <errno.h>
#include <time.h>

#include "cn_did_wifi.h"

#if __BYTE_ORDER__ == __ORDER_LITTLE_ENDIAN__
#define cpu_to_le16(x)  (x)
#define cpu_to_le64(x)  (x)
#else
#define cpu_to_le16(x)      (bswap_16(x))
#define cpu_to_le64(x)      (bswap_64(x))
#endif

#define IEEE80211_FCTL_FTYPE          0x000c
#define IEEE80211_FCTL_STYPE          0x00f0

#define IEEE80211_FTYPE_MGMT            0x0000
#define IEEE80211_STYPE_ACTION          0x00D0
#define IEEE80211_STYPE_BEACON          0x0080

/* IEEE 802.11-2016 capability info */
#define IEEE80211_CAPINFO_ESS               0x0001
#define IEEE80211_CAPINFO_IBSS              0x0002
#define IEEE80211_CAPINFO_CF_POLLABLE       0x0004
#define IEEE80211_CAPINFO_CF_POLLREQ        0x0008
#define IEEE80211_CAPINFO_PRIVACY           0x0010
#define IEEE80211_CAPINFO_SHORT_PREAMBLE    0x0020
/* bits 6-7 reserved */
#define IEEE80211_CAPINFO_SPECTRUM_MGMT     0x0100
#define IEEE80211_CAPINFO_QOS               0x0200
#define IEEE80211_CAPINFO_SHORT_SLOTTIME    0x0400
#define IEEE80211_CAPINFO_APSD              0x0800
#define IEEE80211_CAPINFO_RADIOMEAS         0x1000
/* bit 13 reserved */
#define IEEE80211_CAPINFO_DEL_BLOCK_ACK     0x4000
#define IEEE80211_CAPINFO_IMM_BLOCK_ACK     0x8000

/* IEEE 802.11 Element IDs */
#define IEEE80211_ELEMID_SSID		0x00
#define IEEE80211_ELEMID_RATES		0x01
#define IEEE80211_ELEMID_VENDOR		0xDD

/* 《邻居感知网络规范》v3.1第2.8.2节
* NAN集群ID是一个MAC地址，其取值范围为50-6F-9A-01-00-00至50-6F-9A-01-FF-FF，
* 承载于部分NAN帧的A3字段中。NAN集群ID由发起NAN集群的设备随机选择。
* 然而，《ASTM远程ID规范》v1.1规定，NAN集群ID必须固定为50-6F-9A-01-00-FF这一值。
 */
static const uint8_t *get_nan_cluster_id(void)
{
    static const uint8_t cluster_id[6] = { 0x50, 0x6F, 0x9A, 0x01, 0x00, 0xFF };
    return cluster_id;
}

static int buf_fill_ieee80211_mgmt(uint8_t *buf, size_t *len, size_t buf_size,
                                   const uint16_t subtype,
                                   const uint8_t *dst_addr,
                                   const uint8_t *src_addr,
                                   const uint8_t *bssid)
{
    if (*len + sizeof(struct ieee80211_mgmt) > buf_size)
        return -ENOMEM;

    struct ieee80211_mgmt *mgmt = (struct ieee80211_mgmt *)(buf + *len);
    mgmt->frame_control = (uint16_t) cpu_to_le16(IEEE80211_FTYPE_MGMT | subtype);
    mgmt->duration = cpu_to_le16(0x0000);
    memcpy(mgmt->da, dst_addr, sizeof(mgmt->da));
    memcpy(mgmt->sa, src_addr, sizeof(mgmt->sa));
    memcpy(mgmt->bssid, bssid, sizeof(mgmt->bssid));
    mgmt->seq_ctrl = cpu_to_le16(0x0000);
    *len += sizeof(*mgmt);

    return 0;
}

static int buf_fill_ieee80211_beacon(uint8_t *buf, size_t *len, size_t buf_size, uint16_t interval_tu)
{
    if (*len + sizeof(struct ieee80211_beacon) > buf_size)
        return -ENOMEM;

    struct ieee80211_beacon *beacon = (struct ieee80211_beacon *)(buf + *len);
    struct timespec ts;
    uint64_t mono_us = 0;

#if defined(CLOCK_MONOTONIC)
    clock_gettime(CLOCK_MONOTONIC, &ts);
    mono_us = (uint64_t)((double) ts.tv_sec * 1e6 + (double) ts.tv_nsec * 1e-3);
#elif defined(CLOCK_REALTIME)
    clock_gettime(CLOCK_REALTIME, &ts);
    mono_us = (uint64_t)((double) ts.tv_sec * 1e6 + (double) ts.tv_nsec * 1e-3);
#elif defined(ARDUINO)
#warning "No REALTIME or MONOTONIC clock, using micros()."
    mono_us = micros();
#else
#warning "Unable to set wifi timestamp."
#endif
    beacon->timestamp = cpu_to_le64(mono_us);
    beacon->beacon_interval = cpu_to_le16(interval_tu);
    beacon->capability = cpu_to_le16(IEEE80211_CAPINFO_SHORT_SLOTTIME | IEEE80211_CAPINFO_SHORT_PREAMBLE);
    *len += sizeof(*beacon);

    return 0;
}

/* void* 不能用 可能是编译器版本的问题 改成uint8_t* 了*/
int cndid_message_build_pack(CNDID_UAS_Data *UAS_Data, uint8_t *pack, size_t buflen)
{
    CNDID_MessagePack_data msg_pack;
    CNDID_MessagePack_encoded *msg_pack_enc;
    size_t len;

    /* create a complete message pack 不知道这里为什么不允许void* 都改成对应的指针类别了*/
    msg_pack.SingleMessageSize = CNDID_MESSAGE_SIZE;
    msg_pack.MsgPackSize = 0;
    for (int i = 0; i < CNDID_BASIC_ID_MAX_MESSAGES; i++) {
        if (UAS_Data->BasicIDValid[i]) {
            if (msg_pack.MsgPackSize >= CNDID_PACK_MAX_MESSAGES)
                return -EINVAL;
            if (encodeCNBasicIDMessage((CNDID_BasicID_encoded*)(&msg_pack.Messages[msg_pack.MsgPackSize]), &UAS_Data->BasicID[i]) == CNDID_SUCCESS)
                msg_pack.MsgPackSize++;
        }
    }
    if (UAS_Data->LocationValid) {
        if (msg_pack.MsgPackSize >= CNDID_PACK_MAX_MESSAGES)
            return -EINVAL;
        if (encodeCNLocationMessage((CNDID_Location_encoded*)(&msg_pack.Messages[msg_pack.MsgPackSize]), &UAS_Data->Location) == CNDID_SUCCESS)
            msg_pack.MsgPackSize++;
    }
    
    if (UAS_Data->SelfIDValid) {
        if (msg_pack.MsgPackSize >= CNDID_PACK_MAX_MESSAGES)
            return -EINVAL;
        if (encodeCNSelfIDMessage((CNDID_SelfID_encoded*)(&msg_pack.Messages[msg_pack.MsgPackSize]), &UAS_Data->SelfID) == CNDID_SUCCESS)
            msg_pack.MsgPackSize++;
    }
    if (UAS_Data->SystemValid) {
        if (msg_pack.MsgPackSize >= CNDID_PACK_MAX_MESSAGES)
            return -EINVAL;
        if (encodeCNSystemMessage((CNDID_System_encoded*)(&msg_pack.Messages[msg_pack.MsgPackSize]), &UAS_Data->System) == CNDID_SUCCESS)
            msg_pack.MsgPackSize++;
    }
    /* check that there is at least one message to send. */
    if (msg_pack.MsgPackSize == 0)
        return -EINVAL;

    /* calculate the exact encoded message pack size. */
    len = sizeof(*msg_pack_enc) - (CNDID_PACK_MAX_MESSAGES - msg_pack.MsgPackSize) * CNDID_MESSAGE_SIZE;

    /* check if there is enough space for the message pack. */
    if (len > buflen)
        return -ENOMEM;

    msg_pack_enc = (CNDID_MessagePack_encoded *) pack;
    if (encodeCNMessagePack(msg_pack_enc, &msg_pack) != CNDID_SUCCESS)
        return -1;

    return (int) len;
}
// 里面的服务id可能后面要改 还有一些指向 open droneid的
int cndid_wifi_build_nan_sync_beacon_frame(char *mac, uint8_t *buf, size_t buf_size)
{
    /* Broadcast address */
    uint8_t target_addr[6] = { 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF };
    uint8_t wifi_alliance_oui[3] = { 0x50, 0x6F, 0x9A };
    /* "org.opendroneid.remoteid" hash */
    uint8_t service_id[6] = { 0x88, 0x69, 0x19, 0x9D, 0x92, 0x09 };
    const uint8_t *cluster_id = get_nan_cluster_id();
    struct ieee80211_vendor_specific *vendor;
    struct nan_master_indication_attribute *master_indication_attr;
    struct nan_cluster_attribute *cluster_attr;
    struct nan_service_id_list_attribute *nsila;
    int ret;
    size_t len = 0;

    /* IEEE 802.11 Management Header */
    ret = buf_fill_ieee80211_mgmt(buf, &len, buf_size, IEEE80211_STYPE_BEACON, target_addr, (uint8_t *)mac, cluster_id);
    if (ret <0)
        return ret;

    /* Beacon */
    ret = buf_fill_ieee80211_beacon(buf, &len, buf_size, 0x0200);
    if (ret <0)
        return ret;

    /* Vendor Specific */
    if (len + sizeof(*vendor) > buf_size)
        return -ENOMEM;

    vendor = (struct ieee80211_vendor_specific *)(buf + len);
    memset(vendor, 0, sizeof(*vendor));
    vendor->element_id = IEEE80211_ELEMID_VENDOR;
    vendor->length = 0x22;
    memcpy(vendor->oui, wifi_alliance_oui, sizeof(vendor->oui));
    vendor->oui_type = 0x13; 
    len += sizeof(*vendor);

    /* NAN Master Indication attribute */
    if (len + sizeof(*master_indication_attr) > buf_size)
        return -ENOMEM;

    master_indication_attr = (struct nan_master_indication_attribute *)(buf + len);
    memset(master_indication_attr, 0, sizeof(*master_indication_attr));
    master_indication_attr->header.attribute_id = 0x00;
    master_indication_attr->header.length = cpu_to_le16(0x0002);
    /* Information that is used to indicate a NAN Device’s preference to serve
     * as the role of Master, with a larger value indicating a higher
     * preference. Values 1 and 255 are used for testing purposes only.
     */
    master_indication_attr->master_preference = 0xFE;
    /* Random factor value 0xEA is recommended by the European Standard */
    master_indication_attr->random_factor = 0xEA;
    len += sizeof(*master_indication_attr);

    /* NAN Cluster attribute */
    if (len + sizeof(*cluster_attr) > buf_size)
        return -ENOMEM;

    cluster_attr = (struct nan_cluster_attribute *)(buf + len);
    memset(cluster_attr, 0, sizeof(*cluster_attr));
    cluster_attr->header.attribute_id = 0x1;
    cluster_attr->header.length = cpu_to_le16(0x000D);
    memcpy(cluster_attr->device_mac, mac, sizeof(cluster_attr->device_mac));
    cluster_attr->random_factor = 0xEA;
    cluster_attr->master_preference = 0xFE;
    cluster_attr->hop_count_to_anchor_master = 0x00;
    memset(cluster_attr->anchor_master_beacon_transmission_time, 0, sizeof(cluster_attr->anchor_master_beacon_transmission_time));
    len += sizeof(*cluster_attr);

    /* NAN attributes */
    if (len + sizeof(*nsila) > buf_size)
        return -ENOMEM;

    nsila = (struct nan_service_id_list_attribute *)(buf + len);
    memset(nsila, 0, sizeof(*nsila));
    nsila->header.attribute_id = 0x02;
    nsila->header.length = cpu_to_le16(0x0006);
    memcpy(nsila->service_id, service_id, sizeof(service_id));
    len += sizeof(*nsila);

    return (int) len;
}

int cndid_wifi_build_message_pack_nan_action_frame(CNDID_UAS_Data *UAS_Data, char *mac,
                                                  uint8_t send_counter,
                                                  uint8_t *buf, size_t buf_size)
{
    /* Neighbor Awareness Networking Specification v3.0 in section 2.8.1
     * NAN网络ID要求目标MAC地址为51-6F-9A-01-00-00*/
    uint8_t target_addr[6] = { 0x51, 0x6F, 0x9A, 0x01, 0x00, 0x00 };
    /* "org.opendroneid.remoteid 哈希值 国际通用的*/
    uint8_t service_id[6] = { 0x88, 0x69, 0x19, 0x9D, 0x92, 0x09 };
    uint8_t wifi_alliance_oui[3] = { 0x50, 0x6F, 0x9A };
    const uint8_t *cluster_id = get_nan_cluster_id();
    struct nan_service_discovery *nsd;
    struct nan_service_descriptor_attribute *nsda;
    struct nan_service_descriptor_extension_attribute *nsdea;
    struct CNDID_service_info *si;
    int ret;
    size_t len = 0;

    /* IEEE 802.11 Management Header */
    ret = buf_fill_ieee80211_mgmt(buf, &len, buf_size, IEEE80211_STYPE_ACTION, target_addr, (uint8_t *)mac, cluster_id);
    if (ret <0)
        return ret;

    /* NAN Service Discovery header */
    if (len + sizeof(*nsd) > buf_size)
        return -ENOMEM;

    nsd = (struct nan_service_discovery *)(buf + len);
    memset(nsd, 0, sizeof(*nsd));
    nsd->category = 0x04;               /* IEEE 802.11 Public Action frame */
    nsd->action_code = 0x09;            /* IEEE 802.11 Public Action frame Vendor Specific*/
    memcpy(nsd->oui, wifi_alliance_oui, sizeof(nsd->oui));
    nsd->oui_type = 0x13;               /* Identify Type and version of the NAN */
    len += sizeof(*nsd);

    /* NAN Attribute for Service Descriptor header */
    if (len + sizeof(*nsda) > buf_size)
        return -ENOMEM;

    nsda = (struct nan_service_descriptor_attribute *)(buf + len);
    nsda->header.attribute_id = 0x3;    /* Service Descriptor Attribute type */
    memcpy(nsda->service_id, service_id, sizeof(service_id));
    /* always 1 */
    nsda->instance_id = 0x01;           /* always 1 */
    nsda->requestor_instance_id = 0x00; /* from triggering frame */
    nsda->service_control = 0x10;       /* follow up */
    len += sizeof(*nsda);

    /* CNDID Service Info Attribute header */
    if (len + sizeof(*si) > buf_size)
        return -ENOMEM;

    si = (struct CNDID_service_info *)(buf + len);
    memset(si, 0, sizeof(*si));
    si->message_counter = send_counter;
    len += sizeof(*si);

    ret = cndid_message_build_pack(UAS_Data, buf + len, buf_size - len);
    if (ret < 0)
        return ret;
    len += ret;

    /* set the lengths according to the message pack lengths */
    nsda->service_info_length = sizeof(*si) + ret;
    nsda->header.length = cpu_to_le16(sizeof(*nsda) - sizeof(struct nan_attribute_header) + nsda->service_info_length);

    /* NAN Attribute for Service Descriptor extension header */
    if (len + sizeof(*nsdea) > buf_size)
        return -ENOMEM;

    nsdea = (struct nan_service_descriptor_extension_attribute *)(buf + len);
    nsdea->header.attribute_id = 0xE;
    nsdea->header.length = cpu_to_le16(0x0004);
    nsdea->instance_id = 0x01;
    nsdea->control = cpu_to_le16(0x0200);
    nsdea->service_update_indicator = send_counter;
    len += sizeof(*nsdea);

    return (int) len;
}

int cndid_wifi_build_message_pack_beacon_frame(CNDID_UAS_Data *UAS_Data, char *mac,
                                              const char *SSID, size_t SSID_len,
                                              uint16_t interval_tu, uint8_t send_counter,
                                              uint8_t *buf, size_t buf_size)
{
    /* Broadcast address */
    uint8_t target_addr[6] = { 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF };
    uint8_t asd_stan_oui[3] = { 0xFA, 0x0B, 0xBC }; // 国标也是0xFA0BBC
    /* Mgmt Beacon frame mandatory fields + IE 221 */
    struct ieee80211_ssid *ssid_s;
    struct ieee80211_supported_rates *rates;
    struct ieee80211_vendor_specific *vendor;

    /* Message Pack */
    struct CNDID_service_info *si;

    int ret;
    size_t len = 0;

    /* IEEE 802.11 Management Header */
    ret = buf_fill_ieee80211_mgmt(buf, &len, buf_size, IEEE80211_STYPE_BEACON, target_addr, (uint8_t *)mac, (uint8_t *)mac);
    if (ret <0)
        return ret;

    /* Mandatory Beacon as of 802.11-2016 Part 11 */
    ret = buf_fill_ieee80211_beacon(buf, &len, buf_size, interval_tu);
    if (ret <0)
        return ret;

    /* SSID: 1-32 bytes */
    if (len + sizeof(*ssid_s) > buf_size)
        return -ENOMEM;

    ssid_s = (struct ieee80211_ssid *)(buf + len);
    if(!SSID || (SSID_len ==0) || (SSID_len > 32))
        return -EINVAL;
    ssid_s->element_id = IEEE80211_ELEMID_SSID;
    ssid_s->length = (uint8_t) SSID_len;
    memcpy(ssid_s->ssid, SSID, ssid_s->length);
    len += sizeof(*ssid_s) + SSID_len;

    /* Supported Rates: 1 record at minimum */
    if (len + sizeof(*rates) > buf_size)
        return -ENOMEM;

    rates = (struct ieee80211_supported_rates *)(buf + len);
    rates->element_id = IEEE80211_ELEMID_RATES;
    rates->length = 1; // One rate only
    rates->supported_rates = 0x8C;     // 6 Mbps
    len += sizeof(*rates);

    /* Vendor Specific Information Element (IE 221) */
    if (len + sizeof(*vendor) > buf_size)
        return -ENOMEM;

    vendor = (struct ieee80211_vendor_specific *)(buf + len);
    vendor->element_id = IEEE80211_ELEMID_VENDOR;
    vendor->length = 0x00;  // Length updated at end of function
    memcpy(vendor->oui, asd_stan_oui, sizeof(vendor->oui));
    vendor->oui_type = 0x0D;
    len += sizeof(*vendor);

    /* CNDID Service Info Attribute header */
    if (len + sizeof(*si) > buf_size)
        return -ENOMEM;

    si = (struct CNDID_service_info *)(buf + len);
    memset(si, 0, sizeof(*si));
    si->message_counter = send_counter;
    len += sizeof(*si);

    ret = cndid_message_build_pack(UAS_Data, buf + len, buf_size - len);
    if (ret < 0)
        return ret;
    len += ret;

    /* set the lengths according to the message pack lengths */
    vendor->length = sizeof(vendor->oui) + sizeof(vendor->oui_type) + sizeof(*si) + ret;

    return (int) len;
}

int cndid_message_process_pack(CNDID_UAS_Data *UAS_Data, uint8_t *pack, size_t buflen)
{
    CNDID_MessagePack_encoded *msg_pack_enc = (CNDID_MessagePack_encoded *) pack;
    size_t size = sizeof(*msg_pack_enc) - CNDID_MESSAGE_SIZE * (CNDID_PACK_MAX_MESSAGES - msg_pack_enc->MsgPackSize);
    if (size > buflen)
        return -ENOMEM;

    cndid_initUasData(UAS_Data);

    if (decodeCNMessagePack(UAS_Data, msg_pack_enc) != CNDID_SUCCESS)
        return -1;

    return (int) size;
}

int cndid_wifi_receive_message_pack_nan_action_frame(CNDID_UAS_Data *UAS_Data,
                                                    char *mac, uint8_t *buf, size_t buf_size)
{
    struct ieee80211_mgmt *mgmt; // 802.11管理帧头
    struct nan_service_discovery *nsd; // NAN服务发现帧头
    struct nan_service_descriptor_attribute *nsda; // NAN服务描述符属性
    struct nan_service_descriptor_extension_attribute *nsdea;  // NAN服务描述符扩展属性
    struct CNDID_service_info *si; // 中国远程ID服务信息
    uint8_t target_addr[6] = { 0x51, 0x6F, 0x9A, 0x01, 0x00, 0x00 };  // 远程ID（Remote ID）分配的专用组播地址
    uint8_t wifi_alliance_oui[3] = { 0x50, 0x6F, 0x9A }; // Wi-Fi联盟OUI
    uint8_t service_id[6] = { 0x88, 0x69, 0x19, 0x9D, 0x92, 0x09 }; // 远程ID服务ID 国际通用的 不需要改
    int ret;
    size_t len = 0;

    /* IEEE 802.11 Management Header */
    if (len + sizeof(*mgmt) > buf_size) // 检查缓冲区是否足够容纳一个管理帧头
        return -EINVAL;
    mgmt = (struct ieee80211_mgmt *)(buf + len);
    if ((mgmt->frame_control & cpu_to_le16(IEEE80211_FCTL_FTYPE | IEEE80211_FCTL_STYPE)) !=
        cpu_to_le16(IEEE80211_FTYPE_MGMT | IEEE80211_STYPE_ACTION))
        return -EINVAL;
    if (memcmp(mgmt->da, target_addr, sizeof(mgmt->da)) != 0)
        return -EINVAL;
    memcpy(mac, mgmt->sa, sizeof(mgmt->sa));

    len += sizeof(*mgmt);

    /* NAN Service Discovery header */
    if (len + sizeof(*nsd) > buf_size)
        return -EINVAL;
    nsd = (struct nan_service_discovery *)(buf + len);
    if (nsd->category != 0x04)
        return -EINVAL;
    if (nsd->action_code != 0x09)
        return -EINVAL;
    if (memcmp(nsd->oui, wifi_alliance_oui, sizeof(wifi_alliance_oui)) != 0)
        return -EINVAL;
    if (nsd->oui_type != 0x13)
        return -EINVAL;
    len += sizeof(*nsd);

    /* NAN Attribute for Service Descriptor header */
    if (len + sizeof(*nsda) > buf_size)
        return -EINVAL;
    nsda = (struct nan_service_descriptor_attribute *)(buf + len);
    if (nsda->header.attribute_id != 0x3)
        return -EINVAL;
    if (memcmp(nsda->service_id, service_id, sizeof(service_id)) != 0)
        return -EINVAL;
    if (nsda->instance_id != 0x01)
        return -EINVAL;
    if (nsda->service_control != 0x10)
        return -EINVAL;
    len += sizeof(*nsda);

    si = (struct CNDID_service_info *)(buf + len);
    ret = cndid_message_process_pack(UAS_Data, buf + len + sizeof(*si), buf_size - len - sizeof(*nsdea));
    if (ret < 0)
        return -EINVAL;
    if (nsda->service_info_length != (sizeof(*si) + ret))
        return -EINVAL;
    if (nsda->header.length != (cpu_to_le16(sizeof(*nsda) - sizeof(struct nan_attribute_header) + nsda->service_info_length)))
        return -EINVAL;
    len += sizeof(*si) + ret;

    /* NAN Attribute for Service Descriptor extension header */
    if (len + sizeof(*nsdea) > buf_size)
        return -ENOMEM;
    nsdea = (struct nan_service_descriptor_extension_attribute *)(buf + len);
    if (nsdea->header.attribute_id != 0xE)
        return -EINVAL;
    if (nsdea->header.length != cpu_to_le16(0x0004))
        return -EINVAL;
    if (nsdea->instance_id != 0x01)
        return -EINVAL;
    if (nsdea->control != cpu_to_le16(0x0200))
        return -EINVAL;

    return 0;
}
