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

#include "stellar_priv.h"
#include "packet_inject_main.h"

struct packet_inject_rule rule = {0};

static void usage(char *cmd)
{
    printf("Usage: %s [options]\n\n", cmd);
    printf("Options:\n");
    printf("  -h <ip>       Host IP address\n");
    printf("  -p <port>     Port number\n");
    printf("  -t <type>     Type of manipulation\n");
    printf("                   Options: tcp-rst, tcp-fin, tcp-payload, tcp-payload-fin-rst, udp-payload, ctrl-msg\n");
    printf("  -c <condition> Condition for manipulation\n");
    printf("                   Options: c2s-packet, s2c-packet\n");
    printf("  -n <number>   Number of packets received before injecting action\n\n");
    printf("Example:\n");
    printf("  %s -h 192.168.1.100 -p 8080 -t tcp-payload -c c2s-packet -n 5\n", cmd);
    printf("  %s -h 2001:db8::1   -p 8080 -t tcp-rst     -c s2c-packet -n 10\n", cmd);
    printf("\n");
}

static int parse_cmd(int argc, char **argv)
{
    int opt = 0;
    const char *host = NULL;
    const char *type = NULL;
    const char *condition = NULL;
    while ((opt = getopt(argc, argv, "h:p:t:c:n:")) != -1)
    {
        switch (opt)
        {
        case 'h':
            host = optarg;
            break;
        case 'p':
            rule.port = htons(atoi(optarg));
            break;
        case 't':
            type = optarg;
            break;
        case 'c':
            condition = optarg;
            break;
        case 'n':
            rule.number = atoi(optarg);
            break;
        default:
            usage(argv[0]);
            break;
        }
    }

    if (host)
    {
        if (inet_pton(AF_INET, host, &rule.v4) != 1)
        {
            if (inet_pton(AF_INET6, host, &rule.v6) != 1)
            {
                printf("unable to convert host %s to IPv4 / IPv6\n", host);
                return -1;
            }
            else
            {
                rule.ip_type = 6;
            }
        }
        else
        {
            rule.ip_type = 4;
        }
    }

    if (type == NULL)
    {
        usage(argv[0]);
        printf("invalid type\n");
        return -1;
    }
    else if (strcmp(type, "tcp-rst") == 0)
    {
        rule.inject_type = INJECT_TYPE_TCP_RST;
    }
    else if (strcmp(type, "tcp-fin") == 0)
    {
        rule.inject_type = INJECT_TYPE_TCP_FIN;
    }
    else if (strcmp(type, "tcp-payload") == 0)
    {
        rule.inject_type = INJECT_TYPE_TCP_PAYLOAD;
    }
    else if (strcmp(type, "tcp-payload-fin-rst") == 0)
    {
        rule.inject_type = INJECT_TYPE_TCP_PAYLOAD_FIN_RST;
    }
    else if (strcmp(type, "udp-payload") == 0)
    {
        rule.inject_type = INJECT_TYPE_UDP_PAYLOAD;
    }
    else if (strcmp(type, "ctrl-msg") == 0)
    {
        rule.inject_type = INJECT_TYPE_CTRL_MSG;
    }
    else
    {
        usage(argv[0]);
        printf("invalid type\n");
        return -1;
    }

    if (condition == NULL)
    {
        usage(argv[0]);
        printf("invalid condition\n");
        return -1;
    }
    else if (strcmp(condition, "c2s-packet") == 0)
    {
        rule.direction = AFTER_RECV_C2S_N_PACKET;
    }
    else if (strcmp(condition, "s2c-packet") == 0)
    {
        rule.direction = AFTER_RECV_S2C_N_PACKET;
    }
    else
    {
        usage(argv[0]);
        printf("invalid condition\n");
        return -1;
    }

    if (rule.number <= 0)
    {
        usage(argv[0]);
        printf("invalid count\n");
        return -1;
    }
    printf("%s load inject rule:\n", argv[0]);
    printf("      host      : %s\n", host);
    printf("      port      : %d\n", ntohs(rule.port));
    printf("      type      : %s\n", type);
    printf("      condition : %s\n", condition);
    printf("      count     : %lu\n\n", rule.number);

    return 0;
}

int packet_inject_main(int argc, char **argv)
{
    if (parse_cmd(argc, argv) != 0)
    {
        return -1;
    }

    return stellar_main(argc, argv);
}

int __attribute__((weak)) main(int argc, char **argv)
{
    return packet_inject_main(argc, argv);
}