/* SPDX-License-Identifier: BSD-2-Clause */
/*
 * Copyright (c) 2020 Infineon Technologies AG
 * All rights reserved.
 */

#ifdef HAVE_CONFIG_H
#include "config.h" // IWYU pragma: keep
#endif

#include <errno.h>      // for errno
#include <fcntl.h>      // for open, O_APPEND, O_CREAT, O_NONBLOCK, O_WR...
#include <netinet/in.h> // for htons, htonl
#include <stdint.h>     // for uint32_t, uint16_t, uint8_t, uint64_t
#include <stdlib.h>     // for rand, free, getenv, malloc, srand
#include <string.h>     // for memcpy, strerror, strcmp, memset
#include <time.h>       // for timespec, clock_gettime, CLOCK_MONOTONIC_RAW
#include <unistd.h>     // for close, STDERR_FILENO, STDOUT_FILENO

#include "tcti-pcap-builder.h"
#include "util-io/io.h"    // for write_all
#include "util/aux_util.h" // for UNUSED

#define LOGMODULE tcti
#include "util/log.h" // for LOG_ERROR, LOG_WARNING, LOG_TRACE

#if defined(__FreeBSD__) || defined(__illumos__)
#define CLOCK_MONOTONIC_RAW CLOCK_MONOTONIC
#endif

/* constants used */
#define PCAP_MAJOR                         0x0001
#define PCAP_MINOR                         0x0000
#define PCAP_BLOCK_TYPE_SHB                0x0A0D0D0A
#define PCAP_BLOCK_TYPE_IDB                0x00000001
#define PCAP_BLOCK_TYPE_EPB                0x00000006
#define PCAP_SHB_BYTE_ORDER_MAGIC          0x1A2B3C4D
#define PCAP_SHB_SECTION_LEN_NOT_SPECIFIED 0xFFFFFFFFFFFFFFFFUL
#define PCAP_IDB_LINKTYPE_IPv4             0x00E4
#define PCAP_IDB_SNAP_LEN_NO_LIMIT         0x0000
#define PCAP_EPB_INTERFACE_ID              0x00000000

#define IPv4_VERSION                       0x4
#define IPv4_TOS_BEST_EFFORT               0x00
#define IPv4_ID_UNUSED                     0x0000
#define IPv4_FLAGS_DONT_FRAGMENT           0x2
#define IPv4_FRAGMENT_OFFSET_UNUSED        0x0000
#define IPv4_TIME_TO_LIVE_MAX              0xFF
#define IPv4_PROTOCOL_TCP                  0x06
#define IPv4_CHECKSUM_UNUSED               0x0000
#define TCP_HOST_PORT                      50000
#define TCP_TPM_PORT                       2321 /* required by TPM 2.0 dissector */
#define TCP_FLAGS_ACK                      0x10
#define TCP_WINDOW_SIZE_MAX                0xFFFF
#define TCP_CHECKSUM_UNUSED                0x0000
#define TCP_URGENT_PTR_UNUSED              0x0000

#define SIZEOF_IN_OCTETS(x)                (sizeof(x) / sizeof(uint32_t))
#define TO_MULTIPLE_OF_4_BYTE(x)           ((((x) - 1) / 4 * 4 + 4) * !!(x))

/*
 * complies to pcap-ng (IETF RFC draft-tuexen-opsawg-pcapng-01)
 * https://tools.ietf.org/id/draft-tuexen-opsawg-pcapng-01.html
 *
 * pcap-ng file stucture:
 *
 *  * section header block          (shb)           |<-- file header
 *  * interface statistics block    (idb)           |
 *
 *  * - enhanced packet block       (epb)           |<-- single tpm req. or rsp.
 *    | * header                    (epb_header)    |
 *    | * - ip package                              |
 *    |   | * header                (ip_header)     |
 *    |   | * - tcp package                         |
 *    |   |   | * header            (tcp_header)    |
 *    |   -   - * tpm req. or resp.                 |
 *    - * footer                    (epb_footer)    |
 */

/* section header block */
typedef struct __attribute__((packed)) {
    uint32_t block_type;
    uint32_t block_len;
    uint32_t byte_order_magic;
    uint16_t major_version;
    uint16_t minor_version;
    uint64_t section_len;
    /* options (optional) */
    uint32_t block_len_cp;
} shb;

/* interface description block */
typedef struct __attribute__((packed)) {
    uint32_t block_type;
    uint32_t block_len;
    uint16_t link_type;
    uint16_t reserved;
    uint32_t snap_len;
    /* options (optional) */
    uint32_t block_len_cp;
} idb;

/* enhanced packet block */
typedef struct __attribute__((packed)) {
    uint32_t block_type;
    uint32_t block_len;
    uint32_t interface_id;
    uint32_t timestamp_high;
    uint32_t timestamp_low;
    uint32_t captured_packet_len;
    uint32_t original_packet_len;
} epb_header;

typedef struct __attribute__((packed)) {
    /* options (optional) */
    uint32_t block_len_cp;
} epb_footer;

/* ipv4 packet */
typedef struct __attribute__((packed)) {
    uint8_t  version_header_len;
    uint8_t  type_of_service;
    uint16_t packet_len;
    uint16_t id;
    uint16_t flags;
    uint8_t  time_to_live;
    uint8_t  protocol;
    uint16_t checksum;
    uint32_t source;
    uint32_t destination;
    /* options (optional) */
} ip_header;

/* tcp segment */
typedef struct __attribute__((packed)) {
    uint16_t source_port;
    uint16_t destination_port;
    uint32_t seq_no;
    uint32_t ack_no;
    uint16_t header_len_flags;
    uint16_t window_size;
    uint16_t checksum;
    uint16_t urgent_ptr;
    /* options (optional) */
} tcp_header;

static int pcap_write_section_header_block(pcap_buider_ctx *ctx, void *buf, size_t buf_len);

static int pcap_write_interface_description_block(pcap_buider_ctx *ctx, void *buf, size_t buf_len);

static int pcap_write_enhanced_packet_block(pcap_buider_ctx *ctx,
                                            void            *buf,
                                            size_t           buf_len,
                                            const void      *payload,
                                            size_t           payload_len,
                                            int              direction);

static int pcap_write_ip_packet(pcap_buider_ctx *ctx,
                                void            *buf,
                                size_t           buf_len,
                                const void      *payload,
                                size_t           payload_len,
                                int              direction);

static int pcap_write_tcp_segment(pcap_buider_ctx *ctx,
                                  void            *buf,
                                  size_t           buf_len,
                                  const void      *payload,
                                  size_t           payload_len,
                                  int              direction);

int
pcap_init(pcap_buider_ctx *ctx) {
    char           *filename = getenv(ENV_PCAP_FILE);
    struct timespec time;
    size_t          buf_len;
    size_t          offset;
    size_t          uret;
    int             ret;

    if (filename == NULL) {
        LOG_TRACE(ENV_PCAP_FILE " not set. Using default PCAP file: " DEFAULT_PCAP_FILE);
        filename = DEFAULT_PCAP_FILE;
    }

    clock_gettime(CLOCK_MONOTONIC_RAW, &time);
    srand(time.tv_nsec);

    /* random ips to associate unique connection with this tcti instance */
    ctx->ip_host = (rand() << 16) | rand();
    ctx->ip_tpm = (rand() << 16) | rand();

    /* random sequence numbers */
    ctx->tcp_sequence_no_host = rand();
    ctx->tcp_sequence_no_tpm = rand();

    if (!strcmp(filename, "stdout") || !strcmp(filename, "-")) {
        ctx->fd = STDOUT_FILENO;
    } else if (!strcmp(filename, "stderr")) {
        ctx->fd = STDERR_FILENO;
    } else {
        ctx->fd = open(filename, O_WRONLY | O_APPEND | O_CREAT | O_NONBLOCK, 0644);
        if (ctx->fd < 0) {
            LOG_ERROR("Failed to open file %s: %s", filename, strerror(errno));
            goto error;
        }
    }

    uint8_t buf[sizeof(shb) + sizeof(idb)];
    buf_len = sizeof(buf);
    offset = 0;

    ret = pcap_write_section_header_block(ctx, buf, buf_len);
    if (ret < 0) {
        return ret;
    }
    offset += ret;

    ret = pcap_write_interface_description_block(ctx, buf + offset, buf_len - offset);
    if (ret < 0) {
        return ret;
    }
    offset += ret;

    /* write file header: SHB and IDB (can be written multiple times to same file) */
    uret = write_all(ctx->fd, buf, offset, -1);
    if (uret != offset) {
        LOG_ERROR("Failed to write to file %s: %s", filename, strerror(errno));
        goto error;
    }

    return 0;

error:
    pcap_deinit(ctx);
    return -1;
}

int
pcap_print(pcap_buider_ctx *ctx, const void *payload, size_t payload_len, int direction) {
    size_t pdu_len;
    size_t uret;
    int    ret;

    if (!payload) {
        return -1;
    }

    /* get required buffer size */
    ret = pcap_write_enhanced_packet_block(ctx, NULL, 0, payload, payload_len, direction);
    if (ret < 0) {
        return ret;
    }
    pdu_len = ret;

    uint8_t *buf = malloc(pdu_len);
    if (!buf) {
        LOG_ERROR("Out of memory");
        return -1;
    }

    ret = pcap_write_enhanced_packet_block(ctx, buf, pdu_len, payload, payload_len, direction);
    if (ret < 0) {
        goto cleanup;
    }
    pdu_len = ret;

    uret = write_all(ctx->fd, buf, pdu_len, -1);
    if (uret != pdu_len) {
        LOG_ERROR("Failed to write to file: %s", strerror(errno));
        ret = -1;
        goto cleanup;
    }

    ret = 0;

cleanup:
    free(buf);
    return ret;
}

void
pcap_deinit(pcap_buider_ctx *ctx) {
    int ret;

    if (ctx->fd != STDOUT_FILENO && ctx->fd != STDERR_FILENO) {
        ret = close(ctx->fd);
        if (ret != 0) {
            LOG_WARNING("Failed to close file: %s", strerror(errno));
        }
    }
}

static int
pcap_write_section_header_block(pcap_buider_ctx *ctx, void *buf, size_t buf_len) {
    UNUSED(ctx);

    if (buf) {
        if (buf_len < sizeof(shb)) {
            return -1;
        }

        shb section_header = {
            .block_type = PCAP_BLOCK_TYPE_SHB,
            .block_len = sizeof(shb),
            .byte_order_magic = PCAP_SHB_BYTE_ORDER_MAGIC,
            .major_version = PCAP_MAJOR,
            .minor_version = PCAP_MINOR,
            .section_len = PCAP_SHB_SECTION_LEN_NOT_SPECIFIED,
            .block_len_cp = sizeof(shb),
        };

        memcpy(buf, &section_header, sizeof(shb));
    }

    return sizeof(shb);
}

static int
pcap_write_interface_description_block(pcap_buider_ctx *ctx, void *buf, size_t buf_len) {
    UNUSED(ctx);

    if (buf) {
        if (buf_len < sizeof(idb)) {
            return -1;
        }

        idb interface_description = {
            .block_type = PCAP_BLOCK_TYPE_IDB,
            .block_len = sizeof(idb),
            .link_type = PCAP_IDB_LINKTYPE_IPv4,
            .reserved = 0,
            .snap_len = PCAP_IDB_SNAP_LEN_NO_LIMIT,
            .block_len_cp = sizeof(idb),
        };

        memcpy(buf, &interface_description, sizeof(idb));
    }

    return sizeof(ip_header);
}

static int
pcap_write_enhanced_packet_block(pcap_buider_ctx *ctx,
                                 void            *buf,
                                 size_t           buf_len,
                                 const void      *payload,
                                 size_t           payload_len,
                                 int              direction) {
    UNUSED(ctx);

    size_t          pdu_len, sdu_len, sdu_padded_len;
    struct timespec ts;
    uint64_t        timestamp;
    int             ret;

    ret = clock_gettime(CLOCK_REALTIME, &ts);
    if (ret != 0) {
        LOG_WARNING("Failed to get time: %s", strerror(errno));
        ts.tv_sec = 0;
        ts.tv_nsec = 0;
    }
    timestamp = (uint64_t)ts.tv_sec * 1000000 + ts.tv_nsec / 1000;

    /* get ip packet size */
    sdu_len = pcap_write_ip_packet(ctx, NULL, 0, payload, payload_len, direction);

    /* apply padding (multiple of 4 bytes) */
    sdu_padded_len = TO_MULTIPLE_OF_4_BYTE(sdu_len);

    pdu_len = sizeof(epb_header) + sdu_padded_len + sizeof(epb_footer);

    epb_header header = {
        .block_type = PCAP_BLOCK_TYPE_EPB,
        .block_len = pdu_len,
        .interface_id = PCAP_EPB_INTERFACE_ID,
        .timestamp_high = (timestamp >> 32) & 0xFFFFFFFF,
        .timestamp_low = timestamp & 0xFFFFFFFF,
        .captured_packet_len = sdu_len,
        .original_packet_len = sdu_len,
    };

    epb_footer footer = {
        .block_len_cp = pdu_len,
    };

    if (buf) {
        if (buf_len < pdu_len) {
            return -1;
        }

        memcpy(buf, &header, sizeof(epb_header));
        buf += sizeof(epb_header);
        pcap_write_ip_packet(ctx, buf, sdu_len, payload, payload_len, direction);
        buf += sdu_len;
        memset(buf, 0x00, sdu_padded_len - sdu_len);
        buf += (sdu_padded_len - sdu_len);
        memcpy(buf, &footer, sizeof(epb_footer));
    }

    return (int)pdu_len;
}

static int
pcap_write_ip_packet(pcap_buider_ctx *ctx,
                     void            *buf,
                     size_t           buf_len,
                     const void      *payload,
                     size_t           payload_len,
                     int              direction) {
    size_t pdu_len, sdu_len;

    /* get tcp frame size */
    sdu_len = pcap_write_tcp_segment(ctx, NULL, 0, payload, payload_len, direction);

    pdu_len = sizeof(ip_header) + sdu_len;

    ip_header header = {
        .version_header_len = (IPv4_VERSION << 4) | SIZEOF_IN_OCTETS(ip_header),
        .type_of_service = htons(IPv4_TOS_BEST_EFFORT),
        .packet_len = htons(pdu_len),
        .id = htons(IPv4_ID_UNUSED),
        .flags = htons(IPv4_FLAGS_DONT_FRAGMENT << 13 | IPv4_FRAGMENT_OFFSET_UNUSED),
        .time_to_live = IPv4_TIME_TO_LIVE_MAX,
        .protocol = IPv4_PROTOCOL_TCP,
        .checksum = htons(IPv4_CHECKSUM_UNUSED),
    };

    if (direction == PCAP_DIR_HOST_TO_TPM) {
        header.source = htonl(ctx->ip_host);
        header.destination = htonl(ctx->ip_tpm);
    } else if (direction == PCAP_DIR_TPM_TO_HOST) {
        header.source = htonl(ctx->ip_tpm);
        header.destination = htonl(ctx->ip_host);
    }

    if (buf) {
        if (buf_len < pdu_len) {
            return -1;
        }

        memcpy(buf, &header, sizeof(ip_header));
        buf += sizeof(ip_header);
        pcap_write_tcp_segment(ctx, buf, sdu_len, payload, payload_len, direction);
    }

    return (int)pdu_len;
}

static int
pcap_write_tcp_segment(pcap_buider_ctx *ctx,
                       void            *buf,
                       size_t           buf_len,
                       const void      *payload,
                       size_t           payload_len,
                       int              direction) {
    size_t pdu_len;

    pdu_len = sizeof(tcp_header) + payload_len;

    tcp_header header = {
        .header_len_flags = htons((SIZEOF_IN_OCTETS(tcp_header) << 12) | TCP_FLAGS_ACK),
        .window_size = htons(TCP_WINDOW_SIZE_MAX),
        .checksum = htons(TCP_CHECKSUM_UNUSED),
        .urgent_ptr = htons(TCP_URGENT_PTR_UNUSED),
    };

    if (direction == PCAP_DIR_HOST_TO_TPM) {
        header.source_port = htons(TCP_HOST_PORT);
        header.destination_port = htons(TCP_TPM_PORT);
        header.seq_no = htonl(ctx->tcp_sequence_no_host);
        header.ack_no = htonl(ctx->tcp_sequence_no_tpm);
    } else if (direction == PCAP_DIR_TPM_TO_HOST) {
        header.source_port = htons(TCP_TPM_PORT);
        header.destination_port = htons(TCP_HOST_PORT);
        header.seq_no = htonl(ctx->tcp_sequence_no_tpm);
        header.ack_no = htonl(ctx->tcp_sequence_no_host);
    }

    if (buf) {
        if (buf_len < pdu_len) {
            return -1;
        }
        if (payload_len > 0 && payload == NULL) {
            return -1;
        }

        memcpy(buf, &header, sizeof(tcp_header));
        buf += sizeof(tcp_header);
        memcpy(buf, payload, payload_len);

        if (direction == PCAP_DIR_HOST_TO_TPM) {
            ctx->tcp_sequence_no_host += payload_len;
            /* SYN, FIN flags would increment the seq. no. but are not implemented */
        } else if (direction == PCAP_DIR_TPM_TO_HOST) {
            ctx->tcp_sequence_no_tpm += payload_len;
            /* SYN, FIN flags would increment the seq. no. but are not implemented */
        }
    }

    return (int)pdu_len;
}
