#include <string.h>
#include <netinet/ip.h>
#include <netinet/ip6.h>
#include <netinet/tcp.h>
#include <netinet/ether.h>

/******************************************************************************
 * Struct
 ******************************************************************************/

struct tcp_hdr
{
    uint16_t src_port;
    uint16_t dst_port;
    uint32_t sent_seq;
    uint32_t recv_ack;
    uint8_t data_off;
    uint8_t flags;
    uint16_t rx_window;
    uint16_t checksum;
    uint16_t urgent_ptr;
} __attribute__((__packed__));

struct vlan_hdr
{
    uint16_t vlan_tci;  // Priority (3) + CFI (1) + Identifier Code (12)
    uint16_t eth_proto; // Ethernet type of encapsulated frame
} __attribute__((__packed__));

/******************************************************************************
 *  CheckSum
 ******************************************************************************/

static uint16_t ipv4_checksum(const struct iphdr *ipv4_hdr, uint16_t hdr_len)
{
    uint32_t sum = 0;
    const uint16_t *ip1 = (const uint16_t *)ipv4_hdr;

    while (hdr_len > 1)
    {
        sum += *ip1++;
        if (sum & 0x80000000)
        {
            sum = (sum & 0xFFFF) + (sum >> 16);
        }
        hdr_len -= 2;
    }

    while (sum >> 16)
    {
        sum = (sum & 0xFFFF) + (sum >> 16);
    }

    return (~sum);
}

static uint16_t tcp_checksum_v4(const void *l4_hdr, uint16_t l4_len, struct in_addr *src_addr, struct in_addr *dst_addr)
{
    uint16_t *ip_src = (uint16_t *)src_addr;
    uint16_t *ip_dst = (uint16_t *)dst_addr;
    const uint16_t *buffer = (u_int16_t *)l4_hdr;

    uint32_t sum = 0;
    size_t len = l4_len;

    while (len > 1)
    {
        sum += *buffer++;
        if (sum & 0x80000000)
        {
            sum = (sum & 0xFFFF) + (sum >> 16);
        }
        len -= 2;
    }

    if (len & 1)
    {
        sum += *((uint8_t *)buffer);
    }

    sum += *(ip_src++);
    sum += *ip_src;
    sum += *(ip_dst++);
    sum += *ip_dst;
    sum += htons(IPPROTO_TCP);
    sum += htons(l4_len);

    while (sum >> 16)
    {
        sum = (sum & 0xFFFF) + (sum >> 16);
    }

    return ((uint16_t)(~sum));
}

static uint16_t tcp_checksum_v6(const void *l4_hdr, uint16_t l4_len, struct in6_addr *src_addr, struct in6_addr *dst_addr)
{
    uint16_t *ip_src = (uint16_t *)src_addr;
    uint16_t *ip_dst = (uint16_t *)dst_addr;
    const uint16_t *buffer = (u_int16_t *)l4_hdr;

    uint32_t sum = 0;
    size_t len = l4_len;

    while (len > 1)
    {
        sum += *buffer++;
        if (sum & 0x80000000)
        {
            sum = (sum & 0xFFFF) + (sum >> 16);
        }
        len -= 2;
    }

    if (len & 1)
    {
        sum += *((uint8_t *)buffer);
    }

    for (int i = 0; i < 8; i++)
    {
        sum += *ip_src;
        ip_src++;
    }

    for (int i = 0; i < 8; i++)
    {
        sum += *ip_dst;
        ip_dst++;
    }
    sum += htons(IPPROTO_TCP);
    sum += htons(l4_len);

    while (sum >> 16)
    {
        sum = (sum & 0xFFFF) + (sum >> 16);
    }

    return ((uint16_t)(~sum));
}

/******************************************************************************
 *  TCP header
 ******************************************************************************/

/*
 * TCP Header Format
 *
 *   0                   1                   2                   3
 *   0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1
 *  +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
 *  |          Source Port          |       Destination Port        |
 *  +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
 *  |                        Sequence Number                        |
 *  +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
 *  |                    Acknowledgment Number                      |
 *  +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
 *  |  Data |           |U|A|P|R|S|F|                               |
 *  | Offset| Reserved  |R|C|S|S|Y|I|            Window             |
 *  |       |           |G|K|H|T|N|N|                               |
 *  +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
 *  |           Checksum            |         Urgent Pointer        |
 *  +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
 *  |                    Options                    |    Padding    |
 *  +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
 *  |                             data                              |
 *  +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
 */

int tcp_header_construct(char *buffer, uint16_t src_port, uint16_t dst_port, uint32_t seq, uint32_t ack, u_char flags, uint16_t window, uint16_t urgent, const char *options, uint16_t options_len)
{
    struct tcp_hdr *tcp_hdr = (struct tcp_hdr *)buffer;

    tcp_hdr->src_port = src_port;
    tcp_hdr->dst_port = dst_port;
    tcp_hdr->sent_seq = htonl(seq);
    tcp_hdr->recv_ack = htonl(ack);
    tcp_hdr->data_off = (sizeof(struct tcp_hdr) + options_len) << 2;
    tcp_hdr->flags = flags;
    tcp_hdr->rx_window = htons(window);
    tcp_hdr->checksum = 0;
    tcp_hdr->urgent_ptr = urgent;
    memcpy(buffer + sizeof(struct tcp_hdr), options, options_len);

    return sizeof(struct tcp_hdr) + options_len;
}

/******************************************************************************
 *  IPv4 header
 ******************************************************************************/

/*
 * Internet Header Format
 *
 *   0                   1                   2                   3
 *   0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1
 *  +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
 *  |Version|  IHL  |Type of Service|          Total Length         |
 *  +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
 *  |         Identification        |Flags|      Fragment Offset    |
 *  +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
 *  |  Time to Live |    Protocol   |         Header Checksum       |
 *  +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
 *  |                       Source Address                          |
 *  +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
 *  |                    Destination Address                        |
 *  +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
 *  |                    Options                    |    Padding    |
 *  +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
 */

int ipv4_header_construct(char *buffer, uint16_t carry_layer_len, struct in_addr *src_addr, struct in_addr *dst_addr, u_char tos, uint16_t id, uint16_t frag, u_char ttl, u_char protocol)
{
    struct iphdr *ip_hdr = (struct iphdr *)buffer;

    ip_hdr->version = 4;
    ip_hdr->ihl = 5;
    ip_hdr->tos = tos;
    ip_hdr->tot_len = htons(sizeof(struct iphdr) + carry_layer_len);
    ip_hdr->id = htons(id);
    ip_hdr->frag_off = htons(frag);
    ip_hdr->ttl = ttl;
    ip_hdr->protocol = protocol;
    ip_hdr->check = 0;
    ip_hdr->saddr = *(uint32_t *)src_addr;
    ip_hdr->daddr = *(uint32_t *)dst_addr;

    ip_hdr->check = ipv4_checksum(ip_hdr, sizeof(struct iphdr));

    return sizeof(struct iphdr);
}

/******************************************************************************
 *  IPv6 header
 ******************************************************************************/

/*
 * IPv6 Header Format
 *
 *   0                   1                   2                   3
 *   0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1
 *  +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
 *  |Version| Traffic Class |           Flow Label                  |
 *  +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
 *  |         Payload Length        |  Next Header  |   Hop Limit   |
 *  +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
 *  |                                                               |
 *  +                                                               +
 *  |                                                               |
 *  +                         Source Address                        +
 *  |                                                               |
 *  +                                                               +
 *  |                                                               |
 *  +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
 *  |                                                               |
 *  +                                                               +
 *  |                                                               |
 *  +                      Destination Address                      +
 *  |                                                               |
 *  +                                                               +
 *  |                                                               |
 *  +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
 */

int ipv6_header_construct(char *buffer, uint16_t carry_layer_len, struct in6_addr *src_addr, struct in6_addr *dst_addr, u_char ttl, u_char protocol)
{
    struct ip6_hdr *ip6_hdr = (struct ip6_hdr *)buffer;

    ip6_hdr->ip6_flow = 0;
    ip6_hdr->ip6_plen = htons(carry_layer_len);
    ip6_hdr->ip6_nxt = protocol;
    ip6_hdr->ip6_hlim = ttl;
    ip6_hdr->ip6_vfc &= 0x0F;
    ip6_hdr->ip6_vfc |= (6U << 4U);

    ip6_hdr->ip6_src = *src_addr;
    ip6_hdr->ip6_dst = *dst_addr;

    return sizeof(struct ip6_hdr);
}

/******************************************************************************
 *  Vlan header
 ******************************************************************************/

int vlan_header_construct(char *buffer, uint16_t tci, uint16_t type)
{
    struct vlan_hdr *vlan_hdr = (struct vlan_hdr *)buffer;

    vlan_hdr->vlan_tci = htons(tci);
    vlan_hdr->eth_proto = htons(type);

    return sizeof(struct vlan_hdr);
}

/******************************************************************************
 *  Ether header
 ******************************************************************************/

int ether_header_construct(char *buffer, struct ether_addr *src_mac, struct ether_addr *dst_mac, uint16_t type)
{
    struct ethhdr *eth_hdr = (struct ethhdr *)buffer;

    memcpy(eth_hdr->h_dest, (u_char *)dst_mac, ETHER_ADDR_LEN);
    memcpy(eth_hdr->h_source, (u_char *)src_mac, ETHER_ADDR_LEN);
    eth_hdr->h_proto = htons(type);

    return sizeof(struct ethhdr);
}

/******************************************************************************
 *  Build Ether + IPv4 + TCP Packet
 ******************************************************************************/

int tcp_packet_v4_construct(
    char *buffer, // buffer
    struct ether_addr *src_mac, struct ether_addr *dst_mac, uint16_t vlan_tci, uint16_t l3_protocol, // Ether
    struct in_addr *src_addr, struct in_addr *dst_addr, u_char tos, u_char ttl, uint16_t id, // IPv4
    uint16_t src_port, uint16_t dst_port, uint32_t seq, uint32_t ack, u_char flags, uint16_t window, // TCP Header
    const char *tcp_options, uint16_t tcp_options_len, // TCP Options
    const char *payload, uint16_t payload_len) // Payload
{
    uint16_t length = 0;

    // Ethernet header
    uint16_t eth_protocol = vlan_tci > 0 ? ETH_P_8021Q : l3_protocol;
    length += ether_header_construct(buffer + length, src_mac, dst_mac, eth_protocol);

    // VLAN header
    if (vlan_tci > 0)
    {
        length += vlan_header_construct(buffer + length, vlan_tci, l3_protocol);
    }

    // IPv4 Header
    u_char protocol = IPPROTO_TCP;
    uint16_t frag = IP_DF;
    length += ipv4_header_construct(buffer + length, sizeof(struct tcphdr) + tcp_options_len + payload_len, src_addr, dst_addr, tos, id, frag, ttl, protocol);

    // TCP header and payload
    uint16_t urgent = 0;
    struct tcp_hdr *tcp_hdr = (struct tcp_hdr *)(buffer + length);
    length += tcp_header_construct((char *)tcp_hdr, src_port, dst_port, seq, ack, flags, window, urgent, tcp_options, tcp_options_len);
    memcpy(buffer + length, payload, payload_len);
    length += payload_len;
    tcp_hdr->checksum = tcp_checksum_v4((void *)tcp_hdr, sizeof(struct tcp_hdr) + tcp_options_len + payload_len, src_addr, dst_addr);

    return length;
}

/******************************************************************************
 *  Build Ether + IPv6 + TCP Packet
 ******************************************************************************/

int tcp_packet_v6_construct(
    char *buffer, // buffer
    struct ether_addr *src_mac, struct ether_addr *dst_mac, uint16_t vlan_tci, uint16_t l3_protocol, // Ether
    struct in6_addr *src_addr, struct in6_addr *dst_addr, u_char ttl, // IPv6
    uint16_t src_port, uint16_t dst_port, uint32_t seq, uint32_t ack, u_char flags, uint16_t window, // TCP Header
    const char *tcp_options, uint16_t tcp_options_len, // TCP Options
    const char *payload, uint16_t payload_len) // Payload
{
    uint16_t length = 0;

    // Ethernet header
    uint16_t eth_protocol = vlan_tci > 0 ? ETH_P_8021Q : l3_protocol;
    length += ether_header_construct(buffer + length, src_mac, dst_mac, eth_protocol);

    // VLAN header
    if (vlan_tci > 0)
    {
        length += vlan_header_construct(buffer + length, vlan_tci, l3_protocol);
    }

    // IPv6 Header
    u_char protocol = IPPROTO_TCP;
    length += ipv6_header_construct(buffer + length, sizeof(struct tcphdr) + tcp_options_len + payload_len, src_addr, dst_addr, ttl, protocol);

    // TCP header and payload
    uint16_t urgent = 0;
    struct tcp_hdr *tcp_hdr = (struct tcp_hdr *)(buffer + length);
    length += tcp_header_construct((char *)tcp_hdr, src_port, dst_port, seq, ack, flags, window, urgent, tcp_options, tcp_options_len);
    memcpy(buffer + length, payload, payload_len);
    length += payload_len;
    tcp_hdr->checksum = tcp_checksum_v6((void *)tcp_hdr, sizeof(struct tcp_hdr) + tcp_options_len + payload_len, src_addr, dst_addr);

    return length;
}