#include <stdio.h>
#include <errno.h>
#include <stdlib.h>
#include <string.h>
#include <assert.h>
#include <arpa/inet.h>

#include "toml.h"
#include "stellar/log.h"
#include "stellar/packet.h"
#include "stellar/session.h"
#include "stellar/stellar.h"
#include "stellar/stellar_mq.h"

#pragma GCC diagnostic ignored "-Wunused-parameter"

#define INJTECT_PLUGIN_LOG_ERROR(logger, format, ...) STELLAR_LOG_ERROR(logger, "inject plugin", format, ##__VA_ARGS__)
#define INJTECT_PLUGIN_LOG_INFO(logger, format, ...) STELLAR_LOG_INFO(logger, "inject plugin", format, ##__VA_ARGS__)

enum inject_type
{
    INJECT_TCP_RST = 1,
    INJECT_TCP_FIN = 2,
    INJECT_TCP_PAYLOAD = 3,
    INJECT_TCP_PAYLOAD_FIN_RST = 4,
    INJECT_UDP_PAYLOAD = 5,
    INJECT_CTRL_MSG = 6,
};

struct config
{
    int family; // AF_INET or AF_INET6
    union
    {
        struct sockaddr_in v4;
        struct sockaddr_in6 v6;
    } addr;
    uint16_t port;
    uint64_t number; // inject packet after (C2S/S2C) direction receiving n packets
    enum inject_type type;
    enum flow_direction direction;
};

struct inject_plugin_ctx
{
    struct config config;
    struct stellar *st;
    struct logger *logger;
    int sess_plug_id;
    int tcp_topic_id;
    int udp_topic_id;
};

struct packet_exdata
{
    enum flow_direction flow_dir;

    union
    {
        struct in_addr v4;
        struct in6_addr v6;
    } src_addr, dst_addr;

    uint16_t src_port; // host byte order
    uint16_t dst_port; // host byte order

    uint16_t tcp_payload_len;
    uint32_t tcp_seq; // host byte order
    uint32_t tcp_ack; // host byte order
    uint8_t tcp_flags;

    uint32_t inc_seq;
    uint32_t inc_ack;
};

/******************************************************************************
 * Config
 ******************************************************************************/

static const char *type_to_str(enum inject_type type)
{
    switch (type)
    {
    case INJECT_TCP_RST:
        return "TCP-RST";
    case INJECT_TCP_FIN:
        return "TCP-FIN";
    case INJECT_TCP_PAYLOAD:
        return "TCP-PAYLOAD";
    case INJECT_TCP_PAYLOAD_FIN_RST:
        return "TCP-PAYLOAD-FIN-RST";
    case INJECT_UDP_PAYLOAD:
        return "UDP-PAYLOAD";
    case INJECT_CTRL_MSG:
        return "CTRL-MSG";
    default:
        return "UNKNOWN";
    }
}

static int load_config(struct logger *logger, struct config *config, const char *file)
{
    int ret = -1;
    char errbuf[200];
    const char *ptr;
    FILE *fp = NULL;
    toml_table_t *root = NULL;
    toml_table_t *sub = NULL;
    memset(config, 0, sizeof(struct config));

    fp = fopen(file, "r");
    if (fp == NULL)
    {
        INJTECT_PLUGIN_LOG_ERROR(logger, "open config file %s failed, %s", file, strerror(errno));
        goto error_out;
    }

    root = toml_parse_file(fp, errbuf, sizeof(errbuf));
    if (root == NULL)
    {
        INJTECT_PLUGIN_LOG_ERROR(logger, "parse config file %s failed, %s", file, errbuf);
        goto error_out;
    }

    sub = toml_table_in(root, "packet_inject");
    if (sub == NULL)
    {
        INJTECT_PLUGIN_LOG_ERROR(logger, "config file missing packet_inject section");
        goto error_out;
    }

    ptr = toml_raw_in(sub, "filter_ip");
    if (ptr == NULL)
    {
        INJTECT_PLUGIN_LOG_ERROR(logger, "config file missing packet_inject->filter_ip");
        goto error_out;
    }
    if (strcmp(ptr, "any") == 0)
    {
        config->family = AF_UNSPEC;
    }
    else if (inet_pton(AF_INET, ptr, &config->addr.v4.sin_addr) == 1)
    {
        config->family = AF_INET;
    }
    else if (inet_pton(AF_INET6, ptr, &config->addr.v6.sin6_addr) == 1)
    {
        config->family = AF_INET6;
    }
    else
    {
        INJTECT_PLUGIN_LOG_ERROR(logger, "parse packet_inject->filter_ip failed, invalid ip address: %s", ptr);
        goto error_out;
    }

    ptr = toml_raw_in(sub, "filter_port");
    if (ptr == NULL)
    {
        INJTECT_PLUGIN_LOG_ERROR(logger, "config file missing packet_inject->filter_port");
        goto error_out;
    }
    config->port = atoi(ptr);

    ptr = toml_raw_in(sub, "filter_dir");
    if (ptr == NULL)
    {
        INJTECT_PLUGIN_LOG_ERROR(logger, "config file missing packet_inject->filter_dir");
        goto error_out;
    }
    if (strcmp(ptr, "C2S") == 0)
    {
        config->direction = FLOW_DIRECTION_C2S;
    }
    else if (strcmp(ptr, "S2C") == 0)
    {
        config->direction = FLOW_DIRECTION_S2C;
    }
    else
    {
        INJTECT_PLUGIN_LOG_ERROR(logger, "parse packet_inject->filter_dir failed, invalid direction: %s", ptr);
        goto error_out;
    }

    ptr = toml_raw_in(sub, "filter_pkts");
    if (ptr == NULL)
    {
        INJTECT_PLUGIN_LOG_ERROR(logger, "config file missing packet_inject->filter_pkts");
        goto error_out;
    }
    config->number = atoi(ptr);
    if (config->number == 0)
    {
        INJTECT_PLUGIN_LOG_ERROR(logger, "parse packet_inject->filter_pkts failed, invalid number: %s", ptr);
        goto error_out;
    }

    ptr = toml_raw_in(sub, "inject_type");
    if (ptr == NULL)
    {
        INJTECT_PLUGIN_LOG_ERROR(logger, "config file missing packet_inject->inject_type");
        goto error_out;
    }
    if (strcmp(ptr, "TCP-RST") == 0)
    {
        config->type = INJECT_TCP_RST;
    }
    else if (strcmp(ptr, "TCP-FIN") == 0)
    {
        config->type = INJECT_TCP_FIN;
    }
    else if (strcmp(ptr, "TCP-PAYLOAD") == 0)
    {
        config->type = INJECT_TCP_PAYLOAD;
    }
    else if (strcmp(ptr, "TCP-PAYLOAD-FIN-RST") == 0)
    {
        config->type = INJECT_TCP_PAYLOAD_FIN_RST;
    }
    else if (strcmp(ptr, "UDP-PAYLOAD") == 0)
    {
        config->type = INJECT_UDP_PAYLOAD;
    }
    else if (strcmp(ptr, "CTRL-MSG") == 0)
    {
        config->type = INJECT_CTRL_MSG;
    }
    else
    {
        INJTECT_PLUGIN_LOG_ERROR(logger, "parse packet_inject->inject_type failed, invalid inject type: %s", ptr);
        goto error_out;
    }

    ret = 0;

error_out:
    if (root)
    {
        toml_free(root);
    }

    if (fp)
    {
        fclose(fp);
    }

    return ret;
}

static void print_config(struct logger *logger, const struct config *config)
{
    char addr_str[INET6_ADDRSTRLEN] = {0};

    switch (config->family)
    {
    case AF_INET:
        inet_ntop(AF_INET, &config->addr.v4, addr_str, INET6_ADDRSTRLEN);
        break;
    case AF_INET6:
        inet_ntop(AF_INET6, &config->addr.v6, addr_str, INET6_ADDRSTRLEN);
        break;
    default:
        snprintf(addr_str, INET6_ADDRSTRLEN, "any");
        break;
    }

    INJTECT_PLUGIN_LOG_INFO(logger, "config->filter_ip   : %s", addr_str);
    INJTECT_PLUGIN_LOG_INFO(logger, "config->filter_port : %d", config->port);
    INJTECT_PLUGIN_LOG_INFO(logger, "config->filter_dir  : %s", config->direction == FLOW_DIRECTION_C2S ? "C2S" : "S2C");
    INJTECT_PLUGIN_LOG_INFO(logger, "config->filter_pkts : %lu", config->number);
    INJTECT_PLUGIN_LOG_INFO(logger, "config->inject_type : %s", type_to_str(config->type));
}

/******************************************************************************
 * Utils
 ******************************************************************************/

static inline void packet_exdata_init(const struct packet *pkt, enum flow_direction dir, struct packet_exdata *pkt_exdata)
{
    memset(pkt_exdata, 0, sizeof(struct packet_exdata));

    pkt_exdata->flow_dir = dir;

    int get_inner_addr = 0;
    int count = packet_get_layer_count(pkt);
    for (int i = count - 1; i >= 0; i--)
    {
        const struct layer *layer = packet_get_layer_by_idx(pkt, i);
        switch (layer->proto)
        {
        case LAYER_PROTO_TCP:
            pkt_exdata->src_port = ntohs(layer->hdr.tcp->th_sport);
            pkt_exdata->dst_port = ntohs(layer->hdr.tcp->th_dport);
            pkt_exdata->tcp_seq = ntohl(layer->hdr.tcp->th_seq);
            pkt_exdata->tcp_ack = ntohl(layer->hdr.tcp->th_ack);
            pkt_exdata->tcp_flags = layer->hdr.tcp->th_flags;
            pkt_exdata->tcp_payload_len = packet_get_payload_len(pkt);
            break;
        case LAYER_PROTO_UDP:
            pkt_exdata->src_port = ntohs(layer->hdr.udp->uh_sport);
            pkt_exdata->dst_port = ntohs(layer->hdr.udp->uh_dport);
            break;
        case LAYER_PROTO_IPV4:
            pkt_exdata->src_addr.v4 = layer->hdr.ip4->ip_src;
            pkt_exdata->dst_addr.v4 = layer->hdr.ip4->ip_dst;
            get_inner_addr = 1;
            break;
        case LAYER_PROTO_IPV6:
            pkt_exdata->src_addr.v6 = layer->hdr.ip6->ip6_src;
            pkt_exdata->dst_addr.v6 = layer->hdr.ip6->ip6_dst;
            get_inner_addr = 1;
            break;
        default:
            break;
        }

        if (get_inner_addr)
        {
            break;
        }
    }
}

static inline uint32_t uint32_add(uint32_t seq, uint32_t inc)
{
    if (seq > UINT32_MAX - inc)
    {
        seq = ((uint64_t)seq + (uint64_t)inc) % (4294967296);
    }
    else
    {
        seq += inc;
    }

    return seq;
}

static void build_and_send_udp_packet(struct inject_plugin_ctx *ctx, struct session *sess, struct packet_exdata *pkt_exdata,
                                      enum flow_direction inject_dir, const char *udp_payload, uint16_t udp_payload_len)
{
    const struct packet *origin_pkt = session_get_first_packet(sess, inject_dir);
    if (origin_pkt == NULL)
    {
        INJTECT_PLUGIN_LOG_ERROR(ctx->logger, "build UDP packet failed, %s origin packet is NULL", inject_dir == FLOW_DIRECTION_C2S ? "C2S" : "S2C");
        return;
    }

    struct packet *build_pkt = packet_build_udp(origin_pkt, udp_payload, udp_payload_len);
    if (build_pkt == NULL)
    {
        INJTECT_PLUGIN_LOG_ERROR(ctx->logger, "build UDP packet failed");
        return;
    }

    stellar_send_build_packet(ctx->st, build_pkt);
}

static void build_and_send_tcp_packet(struct inject_plugin_ctx *ctx, struct session *sess, struct packet_exdata *pkt_exdata,
                                      enum flow_direction inject_dir, uint8_t tcp_flags, const char *tcp_payload, uint16_t tcp_payload_len)
{
    uint32_t tcp_seq = 0;
    uint32_t tcp_ack = 0;

    /*
     * +--------+  current packet  +---------+     C2S RST      +--------+
     * |        |----------------->|         |----------------->|        |
     * | Client |                  | Stellar |                  | Server |
     * |        |<-----------------|         |<-----------------|        |
     * +--------+    S2C RST       +---------+                  +--------+
     *
     * for example: current packet is C2S
     *
     * inject direction == current direction (inject C2S RST)
     *      tcp_seq = current_packet_seq
     *      tcp_ack = current_packet_ack
     *
     * inject direction != current direction (inject S2C RST)
     *      tcp_seq = current_packet_ack
     *      tcp_ack = current_packet_seq + current_packet_payload_len
     *  or if current packet is a SYN-ACK packet
     *      tcp_seq = current_packet_seq
     *      tcp_ack = current_packet_ack + current_packet_payload_len + 1
     */

    if (inject_dir == pkt_exdata->flow_dir)
    {
        tcp_seq = uint32_add(pkt_exdata->tcp_seq, pkt_exdata->inc_seq);
        tcp_ack = pkt_exdata->tcp_ack;

        pkt_exdata->inc_seq += tcp_payload_len;
        pkt_exdata->inc_seq += (tcp_flags & TH_FIN) ? 1 : 0; // inject RST packer after FIN packer, tcp_seq should be increased by 1
    }
    else
    {
        tcp_seq = uint32_add(pkt_exdata->tcp_ack, pkt_exdata->inc_ack);
        tcp_ack = uint32_add(pkt_exdata->tcp_seq, pkt_exdata->tcp_payload_len + ((pkt_exdata->tcp_flags & TH_SYN) ? 1 : 0));

        pkt_exdata->inc_ack += tcp_payload_len;
        pkt_exdata->inc_ack += (tcp_flags & TH_FIN) ? 1 : 0; // inject RST packer after FIN packer, ack should be increased by 1
    }

    const struct packet *origin_pkt = session_get_first_packet(sess, inject_dir);
    if (origin_pkt == NULL)
    {
        INJTECT_PLUGIN_LOG_ERROR(ctx->logger, "build TCP packet failed, %s origin packet is NULL", inject_dir == FLOW_DIRECTION_C2S ? "C2S" : "S2C");
        return;
    }

    struct packet *build_pkt = packet_build_tcp(origin_pkt, tcp_seq, tcp_ack, tcp_flags, NULL, 0, tcp_payload, tcp_payload_len);
    if (build_pkt == NULL)
    {
        INJTECT_PLUGIN_LOG_ERROR(ctx->logger, "build TCP packet failed");
        return;
    }

    stellar_send_build_packet(ctx->st, build_pkt);
}

/******************************************************************************
 * Core logic
 ******************************************************************************/

static void *on_sess_new(struct session *sess, void *plugin_ctx)
{
    struct inject_plugin_ctx *ctx = (struct inject_plugin_ctx *)plugin_ctx;
    INJTECT_PLUGIN_LOG_INFO(ctx->logger, "handle session new: %s", session_get0_readable_addr(sess));
    return NULL;
}

static void on_sess_free(struct session *sess, void *sess_ctx, void *plugin_ctx)
{
    struct inject_plugin_ctx *ctx = (struct inject_plugin_ctx *)plugin_ctx;
    INJTECT_PLUGIN_LOG_INFO(ctx->logger, "handle session free: %s", session_get0_readable_addr(sess));
}

static void on_sess_msg(struct session *sess, int topic_id, const void *msg, void *sess_ctx, void *plugin_ctx)
{
    if (msg == NULL)
    {
        return;
    }

    char buffer[1024] = {0};
    struct packet *pkt = (struct packet *)msg;
    struct inject_plugin_ctx *ctx = (struct inject_plugin_ctx *)plugin_ctx;
    struct config *config = &ctx->config;
    enum flow_direction flow_dir = session_get_current_flow_direction(sess);
    INJTECT_PLUGIN_LOG_INFO(ctx->logger, "handle session msg: %s (C2S received packets: %lu, S2C received packets: %lu)",
                            session_get0_readable_addr(sess),
                            session_get_stat(sess, FLOW_DIRECTION_C2S, STAT_RAW_PACKETS_RECEIVED),
                            session_get_stat(sess, FLOW_DIRECTION_S2C, STAT_RAW_PACKETS_RECEIVED));

    struct packet_exdata pkt_exdata;
    packet_exdata_init(pkt, flow_dir, &pkt_exdata);

    if (config->family == AF_INET &&
        memcmp(&config->addr.v4, &pkt_exdata.src_addr.v4, sizeof(struct in_addr)) != 0 &&
        memcmp(&config->addr.v4, &pkt_exdata.dst_addr.v4, sizeof(struct in_addr)) != 0)
    {
        return;
    }

    if (config->family == AF_INET6 &&
        memcmp(&config->addr.v6, &pkt_exdata.src_addr.v6, sizeof(struct in6_addr)) != 0 &&
        memcmp(&config->addr.v6, &pkt_exdata.dst_addr.v6, sizeof(struct in6_addr)) != 0)
    {
        return;
    }

    if (config->port &&
        pkt_exdata.src_port != config->port &&
        pkt_exdata.dst_port != config->port)
    {
        return;
    }

    if (session_get_stat(sess, FLOW_DIRECTION_C2S, STAT_INJECTED_PACKETS_SUCCESS) > 0 ||
        session_get_stat(sess, FLOW_DIRECTION_S2C, STAT_INJECTED_PACKETS_SUCCESS) > 0)
    {
        return;
    }

    if (config->direction == FLOW_DIRECTION_C2S && session_get_stat(sess, FLOW_DIRECTION_C2S, STAT_RAW_PACKETS_RECEIVED) != config->number)
    {
        return;
    }

    if (config->direction == FLOW_DIRECTION_S2C && session_get_stat(sess, FLOW_DIRECTION_S2C, STAT_RAW_PACKETS_RECEIVED) != config->number)
    {
        return;
    }

    switch (config->type)
    {
    case INJECT_TCP_RST:
        build_and_send_tcp_packet(ctx, sess, &pkt_exdata, FLOW_DIRECTION_C2S, TH_RST | TH_ACK, NULL, 0);
        build_and_send_tcp_packet(ctx, sess, &pkt_exdata, FLOW_DIRECTION_S2C, TH_RST | TH_ACK, NULL, 0);
        session_set_discard(sess);
        break;
    case INJECT_TCP_FIN:
        build_and_send_tcp_packet(ctx, sess, &pkt_exdata, FLOW_DIRECTION_C2S, TH_FIN | TH_ACK, NULL, 0);
        build_and_send_tcp_packet(ctx, sess, &pkt_exdata, FLOW_DIRECTION_S2C, TH_FIN | TH_ACK, NULL, 0);
        session_set_discard(sess);
        break;
    case INJECT_TCP_PAYLOAD:
        snprintf(buffer, sizeof(buffer), "HTTP/1.1 200 OK\r\nContent-Length: %d\r\n\r\n%s", 5 + 5 + 2, "Hello");
        build_and_send_tcp_packet(ctx, sess, &pkt_exdata, FLOW_DIRECTION_S2C, TH_ACK, buffer, strlen(buffer)); // inject payload to client
        build_and_send_tcp_packet(ctx, sess, &pkt_exdata, FLOW_DIRECTION_S2C, TH_ACK, "World\r\n", 7);         // inject payload to client
        build_and_send_tcp_packet(ctx, sess, &pkt_exdata, FLOW_DIRECTION_S2C, TH_RST | TH_ACK, NULL, 0);       // inject RST to client
        build_and_send_tcp_packet(ctx, sess, &pkt_exdata, FLOW_DIRECTION_C2S, TH_RST | TH_ACK, NULL, 0);       // inject RST to server
        session_set_discard(sess);
        break;
    case INJECT_TCP_PAYLOAD_FIN_RST:
        snprintf(buffer, sizeof(buffer), "HTTP/1.1 200 OK\r\nContent-Length: %d\r\n\r\n%s", 5 + 5 + 2, "Hello");
        build_and_send_tcp_packet(ctx, sess, &pkt_exdata, FLOW_DIRECTION_S2C, TH_ACK, buffer, strlen(buffer)); // inject payload to client
        build_and_send_tcp_packet(ctx, sess, &pkt_exdata, FLOW_DIRECTION_S2C, TH_ACK, "World\r\n", 7);         // inject payload to client
        build_and_send_tcp_packet(ctx, sess, &pkt_exdata, FLOW_DIRECTION_S2C, TH_FIN | TH_ACK, NULL, 0);       // inject FIN to client
        build_and_send_tcp_packet(ctx, sess, &pkt_exdata, FLOW_DIRECTION_S2C, TH_RST | TH_ACK, NULL, 0);       // inject RST to client
        build_and_send_tcp_packet(ctx, sess, &pkt_exdata, FLOW_DIRECTION_C2S, TH_FIN | TH_ACK, NULL, 0);       // inject FIN to server
        build_and_send_tcp_packet(ctx, sess, &pkt_exdata, FLOW_DIRECTION_C2S, TH_RST | TH_ACK, NULL, 0);       // inject RST to server
        session_set_discard(sess);
        break;
    case INJECT_UDP_PAYLOAD:
        build_and_send_udp_packet(ctx, sess, &pkt_exdata, FLOW_DIRECTION_C2S, "Hello Server", 12);
        build_and_send_udp_packet(ctx, sess, &pkt_exdata, FLOW_DIRECTION_S2C, "Hello Client", 12);
        session_set_discard(sess);
        break;
    case INJECT_CTRL_MSG:
        // TOOD
        break;
    default:
        break;
    }
}

/******************************************************************************
 * Plugin API
 ******************************************************************************/

extern "C"
{
    void *packet_inject_init(struct stellar *st)
    {
        struct inject_plugin_ctx *ctx = (struct inject_plugin_ctx *)calloc(1, sizeof(struct inject_plugin_ctx));
        if (ctx == NULL)
        {
            return NULL;
        }

        ctx->st = st;
        ctx->logger = stellar_get_logger(st);
        INJTECT_PLUGIN_LOG_INFO(ctx->logger, "packet inject plugin init");

        if (load_config(ctx->logger, &ctx->config, "./plugin/inject.toml") == -1)
        {
            INJTECT_PLUGIN_LOG_ERROR(ctx->logger, "load config failed");
            free(ctx);
            return NULL;
        }
        print_config(ctx->logger, &ctx->config);

        ctx->sess_plug_id = stellar_session_plugin_register(st, on_sess_new, on_sess_free, ctx);
        ctx->tcp_topic_id = stellar_mq_get_topic_id(st, TOPIC_TCP);
        ctx->udp_topic_id = stellar_mq_get_topic_id(st, TOPIC_UDP);

        stellar_session_mq_subscribe(st, ctx->tcp_topic_id, on_sess_msg, ctx->sess_plug_id);
        stellar_session_mq_subscribe(st, ctx->udp_topic_id, on_sess_msg, ctx->sess_plug_id);

        return ctx;
    }

    void packet_inject_exit(void *plugin_ctx)
    {
        struct inject_plugin_ctx *ctx = (struct inject_plugin_ctx *)plugin_ctx;
        if (ctx)
        {
            INJTECT_PLUGIN_LOG_INFO(ctx->logger, "packet inject plugin exit");
            free(ctx);
        }
    }
}