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

#include "toml.h"
#include "stellar/stellar.h"
#include "stellar/layer.h"
#include "stellar/session.h"
#include "stellar/stellar_mq.h"


#pragma GCC diagnostic ignored "-Wunused-parameter"

#define LOG_ERR(fmt, ...) printf("ERROR [packet inject] " fmt, ##__VA_ARGS__)
#define LOG_INFO(fmt, ...) printf("INFO [packet inject] " fmt, ##__VA_ARGS__)

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

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;
};

static const char *inject_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 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)
    {
        LOG_ERR("open config file %s failed, %s\n", file, strerror(errno));
        goto error_out;
    }

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

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

    ptr = toml_raw_in(sub, "filter_ip");
    if (ptr == NULL)
    {
        LOG_ERR("config file missing packet_inject->filter_ip\n");
        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
    {
        LOG_ERR("parse packet_inject->filter_ip failed, invalid ip address: %s\n", ptr);
        goto error_out;
    }

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

    ptr = toml_raw_in(sub, "filter_dir");
    if (ptr == NULL)
    {
        LOG_ERR("config file missing packet_inject->filter_dir\n");
        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
    {
        LOG_ERR("parse packet_inject->filter_dir failed, invalid direction: %s\n", ptr);
        goto error_out;
    }

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

    ptr = toml_raw_in(sub, "inject_type");
    if (ptr == NULL)
    {
        LOG_ERR("config file missing packet_inject->inject_type\n");
        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
    {
        LOG_ERR("parse packet_inject->inject_type failed, invalid inject type: %s\n", ptr);
        goto error_out;
    }

    ret = 0;

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

    if (fp)
    {
        fclose(fp);
    }

    return ret;
}

static void print_config(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;
    }

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

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

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;
};

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;
    struct layer layer;
    PACKET_FOREACH_LAYER_REVERSE(pkt, layer)
    {
        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 craft_and_send_udp_packet(struct stellar *st, 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)
    {
        LOG_ERR("craft UDP packet failed, %s origin packet is NULL\n", inject_dir == FLOW_DIRECTION_C2S ? "C2S" : "S2C");
        return;
    }

    struct packet *craft_pkt = craft_udp_packet(origin_pkt, udp_payload, udp_payload_len);
    if (craft_pkt == NULL)
    {
        LOG_ERR("craft UDP packet failed\n");
        return;
    }

    stellar_send_crafted_packet(st, craft_pkt);
}

static void craft_and_send_tcp_packet(struct stellar *st, 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)
    {
        LOG_ERR("craft TCP packet failed, %s origin packet is NULL\n", inject_dir == FLOW_DIRECTION_C2S ? "C2S" : "S2C");
        return;
    }

    struct packet *craft_pkt = craft_tcp_packet(origin_pkt, tcp_seq, tcp_ack, tcp_flags, NULL, 0, tcp_payload, tcp_payload_len);
    if (craft_pkt == NULL)
    {
        LOG_ERR("craft TCP packet failed\n");
        return;
    }

    stellar_send_crafted_packet(st, craft_pkt);
}

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

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

static void *on_sess_new(struct session *sess, void *plugin_ctx)
{
    // struct plugin_ctx *ctx = (struct plugin_ctx *)plugin_ctx;
    LOG_INFO("handle session new: %s\n", session_get0_readable_addr(sess));
    return NULL;
}

static void on_sess_free(struct session *sess, void *sess_ctx, void *plugin_ctx)
{
    // struct plugin_ctx *ctx = (struct plugin_ctx *)plugin_ctx;
    LOG_INFO("handle session free: %s\n", 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;// session closing, return directly

    char buffer[1024] = {0};
    struct packet *pkt = (struct packet *)msg;
    struct plugin_ctx *ctx = (struct plugin_ctx *)plugin_ctx;
    struct stellar *st = ctx->st;
    struct config *config = &ctx->config;
    enum flow_direction flow_dir = session_get_current_flow_direction(sess);
    LOG_INFO("handle session msg: %s (C2S received packets: %lu, S2C received packets: %lu)\n",
             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:
        craft_and_send_tcp_packet(st, sess, &pkt_exdata, FLOW_DIRECTION_C2S, TH_RST | TH_ACK, NULL, 0);
        craft_and_send_tcp_packet(st, sess, &pkt_exdata, FLOW_DIRECTION_S2C, TH_RST | TH_ACK, NULL, 0);
        session_set_discard(sess);
        break;
    case INJECT_TCP_FIN:
        craft_and_send_tcp_packet(st, sess, &pkt_exdata, FLOW_DIRECTION_C2S, TH_FIN | TH_ACK, NULL, 0);
        craft_and_send_tcp_packet(st, 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");
        craft_and_send_tcp_packet(st, sess, &pkt_exdata, FLOW_DIRECTION_S2C, TH_ACK, buffer, strlen(buffer)); // inject payload to client
        craft_and_send_tcp_packet(st, sess, &pkt_exdata, FLOW_DIRECTION_S2C, TH_ACK, "World\r\n", 7);         // inject payload to client
        craft_and_send_tcp_packet(st, sess, &pkt_exdata, FLOW_DIRECTION_S2C, TH_RST | TH_ACK, NULL, 0);       // inject RST to client
        craft_and_send_tcp_packet(st, 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");
        craft_and_send_tcp_packet(st, sess, &pkt_exdata, FLOW_DIRECTION_S2C, TH_ACK, buffer, strlen(buffer)); // inject payload to client
        craft_and_send_tcp_packet(st, sess, &pkt_exdata, FLOW_DIRECTION_S2C, TH_ACK, "World\r\n", 7);         // inject payload to client
        craft_and_send_tcp_packet(st, sess, &pkt_exdata, FLOW_DIRECTION_S2C, TH_FIN | TH_ACK, NULL, 0);       // inject FIN to client
        craft_and_send_tcp_packet(st, sess, &pkt_exdata, FLOW_DIRECTION_S2C, TH_RST | TH_ACK, NULL, 0);       // inject RST to client
        craft_and_send_tcp_packet(st, sess, &pkt_exdata, FLOW_DIRECTION_C2S, TH_FIN | TH_ACK, NULL, 0);       // inject FIN to server
        craft_and_send_tcp_packet(st, 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:
        craft_and_send_udp_packet(st, sess, &pkt_exdata, FLOW_DIRECTION_C2S, "Hello Server", 12);
        craft_and_send_udp_packet(st, 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 plugin_ctx *ctx = (struct plugin_ctx *)calloc(1, sizeof(struct plugin_ctx));
        if (ctx == NULL)
        {
            return NULL;
        }

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

        ctx->st = st;
        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);

        LOG_INFO("init\n");

        return ctx;
    }

    void packet_inject_exit(void *plugin_ctx)
    {
        struct plugin_ctx *ctx = (struct plugin_ctx *)plugin_ctx;
        if (ctx)
        {
            LOG_INFO("exit\n");
            free(ctx);
        }
    }
}