//
// Created by conntrack on 4/30/23.
//


#include "netguard.h"

struct ng_session *debug_socket;


const char* debug_src_ip=""; // Android wlan IP
const char* debug_dest_ip=""; // Debug server pub IP

const uint16_t sport = 40408; // local port
const uint16_t dport = 50508; // server port



// pseudo header needed for tcp header checksum calculation
struct pseudo_header
{
    u_int32_t source_address;
    u_int32_t dest_address;
    u_int8_t placeholder;
    u_int8_t protocol;
    u_int16_t tcp_length;
};

#define DATAGRAM_LEN 4096
#define OPT_SIZE 20

unsigned short checksum(const char *buf, unsigned size)
{
    unsigned sum = 0, i;

    /* Accumulate checksum */
    for (i = 0; i < size - 1; i += 2)
    {
        unsigned short word16 = *(unsigned short *) &buf[i];
        sum += word16;
    }

    /* Handle odd-sized case */
    if (size & 1)
    {
        unsigned short word16 = (unsigned char) buf[i];
        sum += word16;
    }

    /* Fold to get the ones-complement result */
    while (sum >> 16) sum = (sum & 0xFFFF)+(sum >> 16);

    /* Invert to get the negative in ones-complement arithmetic */
    return ~sum;
}


void create_data_packet(char** out_packet, int* out_packet_len, struct tcp_session tcps)
{
    // datagram to represent the packet
    char *datagram = calloc(DATAGRAM_LEN, sizeof(char));

    // required structs for IP and TCP header
    struct iphdr *iph = (struct iphdr*)datagram;
    struct tcphdr *tcph = (struct tcphdr*)(datagram + sizeof(struct iphdr));
    struct pseudo_header psh;

    char source_ip[32];
    struct sockaddr_in sin;

    //some address resolution
    strcpy(source_ip , debug_src_ip); // cli ip
    sin.sin_family = AF_INET;
    sin.sin_port = htons(dport); // server port
    sin.sin_addr.s_addr = inet_addr (debug_dest_ip); // server ip


    // IP header configuration
    iph->ihl = 5;
    iph->version = 4;
    iph->tos = 0;
    iph->tot_len = htons(sizeof(struct iphdr) + sizeof(struct tcphdr) + OPT_SIZE);
    iph->id = htons(54321);
    iph->frag_off = 0;
    iph->ttl = 64;
    iph->protocol = IPPROTO_TCP;
    iph->check = 0; // do calc later
    iph->saddr = inet_addr ( source_ip );
    iph->daddr = sin.sin_addr.s_addr;

    // TCP header configuration
    tcph->source = htons (sport);
    tcph->dest = htons (dport);
    tcph->seq = htonl(rand() % 4294967295);
    tcph->ack_seq = htonl(0);
    tcph->doff = 10; // tcp header size
    tcph->fin = 0;
    tcph->syn = 1;
    tcph->rst = 0;
    tcph->psh = 0;
    tcph->ack = 0;
    tcph->urg = 0;
    tcph->check = 0;
    tcph->window = htons(16000); // window size
    tcph->urg_ptr = 0;


    // TCP pseudo header for checksum calculation
    psh.source_address = inet_addr ( source_ip );
    psh.dest_address = sin.sin_addr.s_addr;
    psh.placeholder = 0;
    psh.protocol = IPPROTO_TCP;
    psh.tcp_length = htons(sizeof(struct tcphdr) + OPT_SIZE);
    int psize = sizeof(struct pseudo_header) + sizeof(struct tcphdr) + OPT_SIZE;
    // fill pseudo packet
    char* pseudogram = malloc(psize);
    memcpy(pseudogram, (char*)&psh, sizeof(struct pseudo_header));
    memcpy(pseudogram + sizeof(struct pseudo_header), tcph, sizeof(struct tcphdr) + OPT_SIZE);

    // TODO: change options to PA
    // TCP options are only set in the SYN packet
    // ---- set mss ----
    datagram[40] = 0x02;
    datagram[41] = 0x04;
    int16_t mss = htons(48); // mss value
    memcpy(datagram + 42, &mss, sizeof(int16_t));
    // ---- enable SACK ----
    datagram[44] = 0x04;
    datagram[45] = 0x02;
    // do the same for the pseudo header
    pseudogram[32] = 0x02;
    pseudogram[33] = 0x04;
    memcpy(pseudogram + 34, &mss, sizeof(int16_t));
    pseudogram[36] = 0x04;
    pseudogram[37] = 0x02;

    tcph->check = checksum((const char*)pseudogram, psize);
    iph->check = checksum((const char*)datagram, iph->tot_len);

    *out_packet = datagram;
    *out_packet_len = sizeof(struct iphdr) + sizeof(struct tcphdr) + OPT_SIZE;
    free(pseudogram);





}

void create_syn_packet(char** out_packet, int* out_packet_len)
{
    // datagram to represent the packet
    char *datagram = calloc(DATAGRAM_LEN, sizeof(char));

    // required structs for IP and TCP header
    struct iphdr *iph = (struct iphdr*)datagram;
    struct tcphdr *tcph = (struct tcphdr*)(datagram + sizeof(struct iphdr));
    struct pseudo_header psh;

    char source_ip[32];
    struct sockaddr_in sin;

    //some address resolution
    strcpy(source_ip , debug_src_ip); // cli ip
    sin.sin_family = AF_INET;
    sin.sin_port = htons(dport); // server port
    sin.sin_addr.s_addr = inet_addr (debug_dest_ip); // server ip


    // IP header configuration
    iph->ihl = 5;
    iph->version = 4;
    iph->tos = 0;
    iph->tot_len = htons(sizeof(struct iphdr) + sizeof(struct tcphdr) + OPT_SIZE);
    iph->id = htons(54321);
    iph->frag_off = 0;
    iph->ttl = 64;
    iph->protocol = IPPROTO_TCP;
    iph->check = 0; // do calc later
    iph->saddr = inet_addr ( source_ip );
    iph->daddr = sin.sin_addr.s_addr;

    // TCP header configuration
    tcph->source = htons (sport);
    tcph->dest = htons (dport);

    tcph->seq = htonl(rand() % 4294967295);
    tcph->ack_seq = htonl(0);
    tcph->doff = 10; // tcp header size
    tcph->fin = 0;
    tcph->syn = 1;
    tcph->rst = 0;
    tcph->psh = 0;
    tcph->ack = 0;
    tcph->urg = 0;
    tcph->check = 0;
    tcph->window = htons(16000); // window size
    tcph->urg_ptr = 0;


    // TCP pseudo header for checksum calculation
    psh.source_address = inet_addr ( source_ip );
    psh.dest_address = sin.sin_addr.s_addr;
    psh.placeholder = 0;
    psh.protocol = IPPROTO_TCP;
    psh.tcp_length = htons(sizeof(struct tcphdr) + OPT_SIZE);
    int psize = sizeof(struct pseudo_header) + sizeof(struct tcphdr) + OPT_SIZE;
    // fill pseudo packet
    char* pseudogram = malloc(psize);
    memcpy(pseudogram, (char*)&psh, sizeof(struct pseudo_header));
    memcpy(pseudogram + sizeof(struct pseudo_header), tcph, sizeof(struct tcphdr) + OPT_SIZE);

    // TCP options are only set in the SYN packet
    // ---- set mss ----
    datagram[40] = 0x02;
    datagram[41] = 0x04;
    int16_t mss = htons(48); // mss value
    memcpy(datagram + 42, &mss, sizeof(int16_t));
    // ---- enable SACK ----
    datagram[44] = 0x04;
    datagram[45] = 0x02;
    // do the same for the pseudo header
    pseudogram[32] = 0x02;
    pseudogram[33] = 0x04;
    memcpy(pseudogram + 34, &mss, sizeof(int16_t));
    pseudogram[36] = 0x04;
    pseudogram[37] = 0x02;

    tcph->check = checksum((const char*)pseudogram, psize);
    iph->check = checksum((const char*)datagram, iph->tot_len);

    *out_packet = datagram;
    *out_packet_len = sizeof(struct iphdr) + sizeof(struct tcphdr) + OPT_SIZE;
    free(pseudogram);

}



int write_data_packet(const struct arguments *args, int epoll_fd, const uint8_t *buffer, size_t length) {
    // send PSH data
    char* psh_packet;
    int psh_packet_len;

    psh_packet = "testoooo";
    psh_packet_len = 8;

    //create_data_packet(&psh_packet, &psh_packet_len, tcps);
    //handle_ip(args, psh_packet, (size_t) psh_packet_len, epoll_fd, 10, 200);

    //write(debug_socket->socket, psh_packet, (size_t) psh_packet_len);
    write(debug_socket->socket, buffer, length);

    //write_ack(args, &debug_socket->tcp); this will send acks from dst to source (wrong direction) if uncommented
    log_android(ANDROID_LOG_ERROR, "Handling push data IP create with length: %d", psh_packet_len);

}




int open_debug_packet(const struct arguments *args, int epoll_fd) {

    // send SYN
    char* packet;
    int packet_len;

    create_syn_packet(&packet, &packet_len);

    handle_ip(args, packet, (size_t) packet_len, epoll_fd, 10, 200);


    /*
    ssize_t res = write(args->tun, packet, (size_t) packet_len);

    if (res >= 0) {
        log_android(ANDROID_LOG_ERROR, "successfuly wrote new syn packet to tun");
        //handle_ip(args, packet, (size_t) packet_len, epoll_fd, 10, 200);
    } else {
        log_android(ANDROID_LOG_ERROR, "tcp write error..");
    }
    */


    return 1;
}




int debug_socket_init(const struct arguments *args, int epoll_fd) {

    log_android(ANDROID_LOG_ERROR, "init debug socket");
    open_debug_packet(args, epoll_fd);

    return 1;

}


struct ng_session *get_debug_session(const struct arguments *args) {


    // Search session
    struct ng_session *cur = args->ctx->ng_session;
    while (cur != NULL &&
           !(cur->protocol == IPPROTO_TCP &&
             cur->tcp.version == 4 &&
             cur->tcp.source == ntohs(40408) && cur->tcp.dest == ntohs(50508)))
        cur = cur->next;


    if (cur == NULL) {
        log_android(ANDROID_LOG_ERROR, "Found null debug session...");
    } else {
        log_android(ANDROID_LOG_ERROR, "Found the debug session..");
        debug_socket = cur;
    }

    return debug_socket;
}


void read_debug_socket() {
    // TODO: Figure out what needs to be passed as parameters to this function
    return ;
}

void write_debug_socket(const struct arguments *args, int epoll_fd, const uint8_t *buffer, size_t length) {
    // TODO: This function is modelled after write_pcap_ret so I made
    // parameters for this function the same since we basically want to do the same thing.

    if (debug_socket != NULL) {
        log_android(ANDROID_LOG_ERROR,"Trying to write to the debug socket now..");
        write_data_packet(args, epoll_fd, buffer, length);
    }


    /*
    struct tcp_session *cur = &debug_socket->tcp;


    // test write to the debug socket
    //write_data(args, cur, buffer, length);

    log_android(ANDROID_LOG_ERROR, "debug tcp port: %d", cur->source);

    int is_debug_server = strcmp(dest_ip, "");
    if (is_debug_server != 0) {

        int res = write_ack(args, &debug_socket->tcp);
        log_android(ANDROID_LOG_ERROR, "write ack result %d", res);
        log_android(ANDROID_LOG_ERROR, "writing debug packet to %s with length: %d", dest_ip, length);

        // Forward to tun
        if (write_data(args, &debug_socket->tcp, buffer, length) >= 0) {
            log_android(ANDROID_LOG_ERROR, "Successfully wrote to debug socket with length: %d", length);
            debug_socket->tcp.local_seq += length;
            debug_socket->tcp.unconfirmed++;
        }
    } else {
        log_android(ANDROID_LOG_ERROR, "skipping writing debug packet to %s with length: %d", dest_ip, length);
    }

    */


}