#include <stdio.h>
#include <stdbool.h>
#include <stdlib.h>
#include <pcap.h>
#include <netinet/if_ether.h>
#include <netinet/ip.h>
#include <netinet/udp.h>
#include <netinet/tcp.h>
#include <netinet/ip_icmp.h>
#include <search.h>

// Global structures
struct scanTypes
{
    u_int connect;
    u_int null;
    u_int halfOpen;
    u_int udp;
    u_int attemptedScans;
    u_int totalPacketsPerAttack;
    u_int xmas;
} scansInFile;

struct scanTypesBools
{
    bool connect;
    bool null;
    bool halfOpen;
    bool udp;
    bool attemptedScans;
    bool xmas;
};

struct tcpFlagsSet
{
    bool ACK;
    bool SYN;
    bool SYN_ACK;
    bool RST_ACK;
    bool FIN;
    bool RST;
    bool PSH;
    bool URG;
};

struct tcpDictionary
{
    struct tcpFlagsSet tcpFlags;
    int packetNumInGroup;
    uint timesSeen;
    struct scanTypesBools scanTypes;
    uint16_t srcPort;
    uint16_t port;
    uint16_t dstPort;
    char *srcIP;
    char *dstIP;
} currentTCPPacket;

struct udpPacket
{
    uint16_t timesSeen;
    uint16_t srcPort;
    uint16_t port;
    uint16_t dstPort;
    char *srcIP;
    char *dstIP;
};

// Provided structures
struct pcap_pkthdr *pcapHeaderStruct;
struct ether_header *ethNetHeaderStruct;
struct ip *ipHeaderStruct;
struct tcphdr *tcpHeaderStruct;
struct udphdr *udpHeaderStruct;
struct icmphdr *icmpHeaderStruct;

// Functions
void myHandler(u_char *args, const struct pcap_pkthdr *header, const u_char *packet);

ENTRY *seachTable(int key);

void addToTable(int key, void *data);

void printScans();

// Global variables and constants
uint *portNums;

int connectScans = 0;

const int TCP_DICTIONARY_SIZE = 100000;
const char *ErrorMSG = "Must provide input file.\n";

ENTRY *search;
ENTRY entry;

int main(int argc, char **argv)
{
    if (hcreate(TCP_DICTIONARY_SIZE) == 0)
    {
        hcreate(50000);
    };
    if (argc < 3)
    {
        printf("%s", ErrorMSG);
        return 0;
    }

    portNums = malloc(TCP_DICTIONARY_SIZE * TCP_DICTIONARY_SIZE * sizeof(uint));

    char errbuf[PCAP_ERRBUF_SIZE];
    char *pcapFileName = argv[2];

    pcap_t *pcapFile = pcap_open_offline(pcapFileName, errbuf);

    pcap_loop(pcapFile, 0, myHandler, NULL);

    /* And close the session */
    pcap_close(pcapFile);

    printScans();
}

void printScans()
{
    printf("Null: %u\n", scansInFile.null);
    printf("Xmas: %u\n", scansInFile.xmas);
    printf("UDP: %u\n", scansInFile.udp);
    printf("Half-open: %u\n", scansInFile.halfOpen);
    printf("Connect: %u\n", scansInFile.connect);
    printf("Attempted scans: %u\n", scansInFile.attemptedScans);
    printf("Total packets per attack: %u\n", scansInFile.totalPacketsPerAttack);
}

void myHandler(
    u_char *args,
    const struct pcap_pkthdr *header,
    const u_char *packet)
{
    /* First, lets make sure we have an IP packet */
    struct ether_header *eth_header;
    eth_header = (struct ether_header *)packet;
    if (ntohs(eth_header->ether_type) != ETHERTYPE_IP)
    {
        return;
    }

    const u_char *ip_header;
    const u_char *tcp_header;
    const u_char *udp_header;
    const u_char *icmp_header;
    const u_char *payload;

    // Packet lengths
    int ethHeaderLength = 14;
    int ipHeaderLength;
    int tcpHeaderLength;
    int payloadLength;

    ip_header = packet + ethHeaderLength;

    ipHeaderLength = ((*ip_header) & 0x0F);

    ipHeaderLength = ipHeaderLength * 4;
    ipHeaderStruct = (struct iphdr *)ip_header;
    // printf("%d\n", ipHeaderStruct->ip_id);

    u_char protocol = *(ip_header + 9);

    struct udpPacket currentUDPPacket;

    /*
        UDP scan is not successful if ICMP is returned
    */
    if (protocol == IPPROTO_ICMP)
    {
        search = seachTable(currentUDPPacket.port);

        struct udpPacket *prevUDPPacket;
        if (search != NULL)
            prevUDPPacket = search->data;
        icmp_header = packet + ethHeaderLength + ipHeaderLength;
        icmpHeaderStruct = (struct icmphdr *)icmp_header;
        if (icmpHeaderStruct->type == ICMP_DEST_UNREACH)
        {
            scansInFile.udp--;
            scansInFile.attemptedScans++;
        }
    }

    // UDP scan is successful if no response from server or response.
    if (protocol == IPPROTO_UDP)
    {
        struct udpPacket *prevUDPPacket;

        udp_header = packet + ethHeaderLength + ipHeaderLength;

        udpHeaderStruct = (struct udphdr *)udp_header;

        currentUDPPacket.port = ntohs(udpHeaderStruct->source);

        search = seachTable(currentUDPPacket.port);

        if (search != NULL)
            prevUDPPacket = search->data;

        currentUDPPacket.dstPort = ntohs(udpHeaderStruct->dest);
        currentUDPPacket.srcPort = ntohs(udpHeaderStruct->source);

        if (currentUDPPacket.dstPort == 53 || currentUDPPacket.srcPort == 53)
        {
            return;
        }

        // response from server
        if (search != NULL && currentUDPPacket.dstPort == prevUDPPacket->port)
        {
            scansInFile.totalPacketsPerAttack++;
            scansInFile.attemptedScans++;
            scansInFile.udp++;
        }

        currentUDPPacket.timesSeen++;

        if (search != NULL && prevUDPPacket->timesSeen > 1)
        {
            scansInFile.udp--;
            addToTable(currentUDPPacket.port, &currentUDPPacket);
            return;
        }

        addToTable(currentUDPPacket.port, &currentUDPPacket);

        scansInFile.udp++;
        return;
    }

    if (protocol != IPPROTO_TCP)
    {
        return;
    }

    tcp_header = packet + ethHeaderLength + ipHeaderLength;

    tcpHeaderLength = ((*(tcp_header + 12)) & 0xF0) >> 4;

    tcpHeaderLength = tcpHeaderLength * 4;
    tcpHeaderStruct = (struct tcphdr *)tcp_header;

    const bool ACKflag = (tcpHeaderStruct->ack == 1);
    const bool RSTflag = (tcpHeaderStruct->rst == 1);
    const bool SYNflag = (tcpHeaderStruct->syn == 1);
    const bool FINflag = (tcpHeaderStruct->fin == 1);
    const bool PSHflag = (tcpHeaderStruct->fin == 1);
    const bool URGflag = (tcpHeaderStruct->urg == 1);
    const uint32_t tcpSEQ = tcpHeaderStruct->seq;
    const uint32_t tcpACKSeq = tcpHeaderStruct->ack_seq;

    currentTCPPacket.tcpFlags.ACK = ACKflag;
    currentTCPPacket.tcpFlags.RST = RSTflag;
    currentTCPPacket.tcpFlags.SYN = SYNflag;
    currentTCPPacket.tcpFlags.FIN = FINflag;
    currentTCPPacket.tcpFlags.PSH = PSHflag;
    currentTCPPacket.tcpFlags.URG = URGflag;

    /*
    TCP SYN flag is set when connection is made and when server sends back SYN/ACK

    ACK flag set when server sends back SYN/ACK, and when client sends response to SYN/ACK
    */

    const uint16_t dstPort = ntohs(tcpHeaderStruct->dest);

    const uint16_t srcPort = ntohs(tcpHeaderStruct->source);
    char *srcIP = inet_ntoa(ipHeaderStruct->ip_dst);

    currentTCPPacket.packetNumInGroup = 0;

    portNums[currentTCPPacket.port]++;
    // start of packet for Half-open and Connect
    if (SYNflag && !ACKflag && !RSTflag && !FINflag && !PSHflag && !URGflag)
    {
        currentTCPPacket.port = dstPort;
        currentTCPPacket.srcPort = srcPort;
        currentTCPPacket.dstPort = dstPort;
        currentTCPPacket.packetNumInGroup = 1;
        portNums[dstPort] = 1;

        addToTable(currentTCPPacket.dstPort, &currentTCPPacket);
        return;
    }
    else if (FINflag)
    {
        // Xmas pattern
        if (PSHflag && URGflag)
        {
            scansInFile.xmas++;
            currentTCPPacket.port = dstPort;
            currentTCPPacket.packetNumInGroup = 1;
            addToTable(currentTCPPacket.port, &currentTCPPacket);
            return;
        }
    }
    // NULL pattern
    if (!SYNflag && !ACKflag &&
        !RSTflag && !FINflag &&
        !PSHflag && !URGflag)
    {
        scansInFile.null++;
        scansInFile.totalPacketsPerAttack++;
        // printf("\n\nDST port: %d\n", dstPort);
        currentTCPPacket.port = dstPort;
        currentTCPPacket.packetNumInGroup = 1;
        addToTable(currentTCPPacket.port, &currentTCPPacket);
        return;
    }

    struct tcpDictionary *prevPacket;

    // SYN, ACK
    if (ACKflag && SYNflag)
    {
        for (uint i = 0; i < TCP_DICTIONARY_SIZE; i++)
        {
            search = seachTable(i);

            if (dstPort == 1029 && search != NULL)
            {
                // printf("\nSYN ACK flag");
                // printf("\n\nDST port: %d\n", currentTCPPacket.dstPort);
                // printf("\nSRC port: %d\n", currentTCPPacket.srcPort);
            }
            if (search != NULL && i == dstPort)
            {
                prevPacket = search->data;

                scansInFile.totalPacketsPerAttack++;
                const uint16_t dstPort = ntohs(tcpHeaderStruct->dest);

                const uint16_t srcPort = ntohs(tcpHeaderStruct->source);
                char *srcIP = inet_ntoa(ipHeaderStruct->ip_dst);

                // does last packet match X-Mas pattern?
                if (prevPacket->tcpFlags.PSH &&
                    prevPacket->tcpFlags.FIN &&
                    prevPacket->tcpFlags.URG)
                {
                    if (FINflag)
                    {
                        // Xmas pattern
                        if (PSHflag && URGflag)
                        {
                            scansInFile.xmas++;
                            currentTCPPacket.port = dstPort;
                            currentTCPPacket.packetNumInGroup = 1;
                            addToTable(currentTCPPacket.port, &currentTCPPacket);
                            return;
                        }
                    }
                    scansInFile.xmas--;
                    scansInFile.attemptedScans++;
                    currentTCPPacket.packetNumInGroup = 0;
                    return;
                }
                else if (prevPacket->tcpFlags.SYN)
                {
                    // printf("SYN, ACK Port %u\n", currentTCPPacket.dstPort);
                    currentTCPPacket.packetNumInGroup++;
                    scansInFile.totalPacketsPerAttack++;
                    currentTCPPacket.port = dstPort;
                    addToTable(currentTCPPacket.port, &currentTCPPacket);
                }
            }
        }
        return;
    }

    // RST, ACK
    else if (ACKflag && RSTflag)
    {
        for (uint i = 0; i < TCP_DICTIONARY_SIZE; i++)
        {
            search = seachTable(i);

            if (search != NULL && i == dstPort)
            {
                prevPacket = search->data;

                // printf("\nRST ACK\n");
                // printf("\nDST port: %d\n", currentTCPPacket.dstPort);
                // printf("\nSRC port: %d\n", currentTCPPacket.srcPort);
                const uint16_t dstPort = ntohs(tcpHeaderStruct->dest);

                const uint16_t srcPort = ntohs(tcpHeaderStruct->source);
                char *srcIP = inet_ntoa(ipHeaderStruct->ip_dst);

                // does last packet match X-Mas pattern?
                if (prevPacket->tcpFlags.PSH &&
                    prevPacket->tcpFlags.FIN &&
                    prevPacket->tcpFlags.URG)
                {
                    if (FINflag)
                    {
                        // Xmas pattern
                        if (PSHflag && URGflag)
                        {
                            portNums[dstPort] = dstPort;
                            scansInFile.xmas++;
                            currentTCPPacket.port = dstPort;
                            currentTCPPacket.packetNumInGroup = 1;
                            addToTable(currentTCPPacket.port, &currentTCPPacket);
                            return;
                        }
                    }
                    scansInFile.xmas--;
                    scansInFile.attemptedScans++;
                    scansInFile.totalPacketsPerAttack++;
                    currentTCPPacket.packetNumInGroup = 0;
                    return;
                }
                // connect scan complete
                else if (prevPacket->tcpFlags.ACK && prevPacket->dstPort == i)
                {
                    portNums[dstPort] = dstPort;
                    scansInFile.connect++;
                    scansInFile.totalPacketsPerAttack++;
                    return;
                }
                if (!prevPacket->tcpFlags.SYN && !prevPacket->tcpFlags.ACK &&
                    !prevPacket->tcpFlags.RST && !prevPacket->tcpFlags.FIN &&
                    !prevPacket->tcpFlags.PSH && !prevPacket->tcpFlags.URG)
                {
                    scansInFile.totalPacketsPerAttack++;
                    scansInFile.null--;
                }
            }
        }
        return;
    }
    // ACK flag
    else if (ACKflag && !SYNflag &&
             !PSHflag && !FINflag &&
             !RSTflag && !URGflag)
    {
        for (uint i = 0; i < TCP_DICTIONARY_SIZE; i++)
        {
            search = seachTable(i);
            if (currentTCPPacket.dstPort == 1029 && search != NULL)
            {

                // printf("\nACK flag");
                // printf("\n\nDST port: %d\n", currentTCPPacket.dstPort);
                // printf("\nSRC port: %d\n", currentTCPPacket.srcPort);
            }
            if (search != NULL && i == dstPort)
            {
                printf("\nFound port: %d\n", currentTCPPacket.srcPort);

                portNums[dstPort]++;
                prevPacket = search->data;
                const uint16_t dstPort = ntohs(tcpHeaderStruct->dest);

                const uint16_t srcPort = ntohs(tcpHeaderStruct->source);
                char *srcIP = inet_ntoa(ipHeaderStruct->ip_dst);

                if (prevPacket->tcpFlags.ACK &&
                    prevPacket->tcpFlags.SYN)
                {
                    scansInFile.totalPacketsPerAttack++;
                    addToTable(currentTCPPacket.dstPort, &currentTCPPacket);
                    return;
                }
                // addToTable(currentTCPPacket.dstPort, &currentTCPPacket);
            }
            return;
        }
    }

    else if (RSTflag)
    {
        for (uint i = 0; i < TCP_DICTIONARY_SIZE; i++)
        {
            search = seachTable(i);

            // printf("\nDST port: %d\n", currentTCPPacket.dstPort);
            // printf("\nSRC port: %d\n", currentTCPPacket.srcPort);
            if (search != NULL)
            {
                // printf("\nFound port: %d\n", currentTCPPacket.port);

                portNums[dstPort]++;
                prevPacket = search->data;
                const uint16_t dstPort = ntohs(tcpHeaderStruct->dest);

                const uint16_t srcPort = ntohs(tcpHeaderStruct->source);
                char *srcIP = inet_ntoa(ipHeaderStruct->ip_dst);

                if (prevPacket->tcpFlags.SYN && prevPacket->tcpFlags.ACK)
                {
                    portNums[dstPort] = dstPort;
                    scansInFile.halfOpen++;
                    scansInFile.totalPacketsPerAttack++;
                    return;
                }
            }
        }
    }
}

void addToTable(int key, void *data)
{
    ENTRY dictionaryEntry;
    char index[50];
    sprintf(index, "%u", key);
    dictionaryEntry.key = index;
    dictionaryEntry.data = data;
    hsearch(dictionaryEntry, ENTER);
}

ENTRY *seachTable(int key)
{
    ENTRY dictionaryEntry;
    char index[50];
    sprintf(index, "%u", key);
    dictionaryEntry.key = index;
    ENTRY *search = hsearch(dictionaryEntry, FIND);
    return search;
}