#include <stdio.h>
#include <stdlib.h>
#include <string.h>

#include "stellar/layer.h"
#include "stellar/session_mq.h"
#include "packet_inject_main.h"

struct packet_inject_plugin_ctx
{
    struct stellar *st;
    int sess_plug_id;
    int tcp_topic_id;
    int udp_topic_id;
    char name[64];
};

static void *on_sess_new(struct session *sess, void *plugin_ctx)
{
    struct packet_inject_plugin_ctx *ctx = (struct packet_inject_plugin_ctx *)plugin_ctx;
    printf("[%s] pluign handle session new: %s\n", ctx->name, session_get0_readable_addr(sess));
    return NULL;
}

static void on_sess_free(struct session *sess, void *sess_ctx, void *plugin_ctx)
{
    struct packet_inject_plugin_ctx *ctx = (struct packet_inject_plugin_ctx *)plugin_ctx;
    printf("[%s] pluign handle session free: %s\n", ctx->name, 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)
{
    struct packet_inject_plugin_ctx *ctx = (struct packet_inject_plugin_ctx *)plugin_ctx;
    printf("[%s] pluign handle session msg: %s (C2S received packets: %lu, S2C received packets: %lu)\n",
           ctx->name, 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 *pkt = (struct packet *)msg;
    char buffer[1024] = {0};
    int is_ip_hit = 0;
    int is_port_hit = 0;
    struct layer layer;
    PACKET_FOREACH_LAYER_REVERSE(pkt, layer)
    {
        switch (layer.proto)
        {
        case LAYER_PROTO_IPV4:
            if (memcmp(&layer.hdr.ip4->ip_src, &rule.addr4, sizeof(struct in_addr)) == 0 ||
                memcmp(&layer.hdr.ip4->ip_dst, &rule.addr4, sizeof(struct in_addr)) == 0)
            {
                is_ip_hit = 1;
            }
            break;
        case LAYER_PROTO_IPV6:
            if (memcmp(&layer.hdr.ip6->ip6_src, &rule.addr6, sizeof(struct in6_addr)) == 0 ||
                memcmp(&layer.hdr.ip6->ip6_dst, &rule.addr6, sizeof(struct in6_addr)) == 0)
            {
                is_ip_hit = 1;
            }
            break;
        case LAYER_PROTO_TCP:
            if (layer.hdr.tcp->th_sport == rule.port ||
                layer.hdr.tcp->th_dport == rule.port)
            {
                is_port_hit = 1;
            }
            break;
        case LAYER_PROTO_UDP:
            if (layer.hdr.udp->uh_sport == rule.port ||
                layer.hdr.udp->uh_dport == rule.port)
            {
                is_port_hit = 1;
            }
            break;
        default:
            break;
        }
    }

    if (rule.family && !is_ip_hit)
    {
        return;
    }
    if (rule.port && !is_port_hit)
    {
        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 (rule.direction == AFTER_RECV_C2S_N_PACKET && session_get_stat(sess, FLOW_DIRECTION_C2S, STAT_RAW_PACKETS_RECEIVED) != rule.number)
    {
        return;
    }
    if (rule.direction == AFTER_RECV_S2C_N_PACKET && session_get_stat(sess, FLOW_DIRECTION_S2C, STAT_RAW_PACKETS_RECEIVED) != rule.number)
    {
        return;
    }
    switch (rule.inject_type)
    {
    case INJECT_TYPE_TCP_RST:
        stellar_inject_tcp_rst(ctx->st, sess, FLOW_DIRECTION_C2S);
        stellar_inject_tcp_rst(ctx->st, sess, FLOW_DIRECTION_S2C);
        session_set_discard(sess);
        break;
    case INJECT_TYPE_TCP_FIN:
        stellar_inject_tcp_fin(ctx->st, sess, FLOW_DIRECTION_C2S);
        stellar_inject_tcp_fin(ctx->st, sess, FLOW_DIRECTION_S2C);
        session_set_discard(sess);
        break;
    case INJECT_TYPE_TCP_PAYLOAD:
        snprintf(buffer, sizeof(buffer), "HTTP/1.1 200 OK\r\nContent-Length: %d\r\n\r\n%s", 5 + 5 + 2, "Hello");
        stellar_inject_tcp_payload(ctx->st, sess, FLOW_DIRECTION_S2C, buffer, strlen(buffer)); // inject payload to client
        stellar_inject_tcp_payload(ctx->st, sess, FLOW_DIRECTION_S2C, "World\r\n", 7);         // inject payload to client
        stellar_inject_tcp_rst(ctx->st, sess, FLOW_DIRECTION_S2C);                             // inject RST to client
        stellar_inject_tcp_rst(ctx->st, sess, FLOW_DIRECTION_C2S);                             // inject RST to server
        session_set_discard(sess);
        break;
    case INJECT_TYPE_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");
        stellar_inject_tcp_payload(ctx->st, sess, FLOW_DIRECTION_S2C, buffer, strlen(buffer)); // inject payload to client
        stellar_inject_tcp_payload(ctx->st, sess, FLOW_DIRECTION_S2C, "World\r\n", 7);         // inject payload to client
        stellar_inject_tcp_fin(ctx->st, sess, FLOW_DIRECTION_S2C);                             // inject FIN to client
        stellar_inject_tcp_rst(ctx->st, sess, FLOW_DIRECTION_S2C);                             // inject RST to client
        stellar_inject_tcp_fin(ctx->st, sess, FLOW_DIRECTION_C2S);                             // inject FIN to server
        stellar_inject_tcp_rst(ctx->st, sess, FLOW_DIRECTION_C2S);                             // inject RST to server
        session_set_discard(sess);
        break;
    case INJECT_TYPE_UDP_PAYLOAD:
        stellar_inject_udp_payload(ctx->st, sess, FLOW_DIRECTION_C2S, "Hello Server", 12);
        stellar_inject_udp_payload(ctx->st, sess, FLOW_DIRECTION_S2C, "Hello Client", 12);
        session_set_discard(sess);
        break;
    case INJECT_TYPE_CTRL_MSG:
        // TOOD
        break;
    default:
        break;
    }
}

extern "C"
{
    void *packet_inject_plugin_init(struct stellar *st)
    {
        struct packet_inject_plugin_ctx *ctx = (struct packet_inject_plugin_ctx *)calloc(1, sizeof(struct packet_inject_plugin_ctx));
        if (ctx == NULL)
        {
            return NULL;
        }
        ctx->st = st;
        ctx->sess_plug_id = stellar_session_plugin_register(st, on_sess_new, on_sess_free, ctx);
        ctx->tcp_topic_id = stellar_session_mq_get_topic_id(st, TOPIC_TCP);
        ctx->udp_topic_id = stellar_session_mq_get_topic_id(st, TOPIC_UDP);
        snprintf(ctx->name, sizeof(ctx->name), "packet_inject");

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

        printf("[%s] plugin init\n", ctx->name);

        return ctx;
    }

    void packet_inject_plugin_exit(void *plugin_ctx)
    {
        struct packet_inject_plugin_ctx *ctx = (struct packet_inject_plugin_ctx *)plugin_ctx;
        if (ctx)
        {
            printf("[%s] plugin exit\n", ctx->name);
            free(ctx);
        }
    }
}