#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <pcap.h>
#include <ndpi_api.h>
#include <netinet/in.h>
#include <netinet/ip.h>
#include <netinet/tcp.h>
#include <netinet/udp.h>

struct ndpi_detection_module_struct *ndpi_struct = NULL;
struct ndpi_flow_struct *ndpi_flow = NULL;

void init_ndpi() {
    struct ndpi_global_context *g_ctx = ndpi_global_init();
    
    ndpi_struct = ndpi_init_detection_module(g_ctx);
    if (ndpi_struct == NULL) {
        printf("ERROR: ndpi_init_detection_module failed\n");
        exit(1);
    }
    
    ndpi_finalize_initialization(ndpi_struct);
    
    // Allocate flow structure  
    ndpi_flow = calloc(1, ndpi_detection_get_sizeof_ndpi_flow_struct());
}

void cleanup_ndpi() {
    if (ndpi_flow) free(ndpi_flow);
    if (ndpi_struct) ndpi_exit_detection_module(ndpi_struct);
}

void analyze_packet_from_hex(const char* hex_data) {
    // Convert hex string to binary data
    size_t hex_len = strlen(hex_data);
    if (hex_len % 2 != 0) {
        printf("ERROR: Invalid hex data length\n");
        return;
    }
    
    size_t bin_len = hex_len / 2;
    unsigned char *packet_data = malloc(bin_len);
    
    for (size_t i = 0; i < bin_len; i++) {
        sscanf(hex_data + 2*i, "%2hhx", &packet_data[i]);
    }
    
    // Basic IP header parsing
    struct iphdr *ip_header = (struct iphdr*)packet_data;
    if (bin_len < sizeof(struct iphdr) || ip_header->version != 4) {
        printf("ERROR: Not a valid IPv4 packet\n");
        free(packet_data);
        return;
    }
    
    uint32_t src_ip = ntohl(ip_header->saddr);
    uint32_t dst_ip = ntohl(ip_header->daddr);
    uint16_t src_port = 0, dst_port = 0;
    uint8_t protocol_type = ip_header->protocol;
    
    // Extract ports and payload
    unsigned char *payload = NULL;
    size_t payload_len = 0;
    size_t header_len = ip_header->ihl * 4;
    
    if (protocol_type == IPPROTO_TCP) {
        if (bin_len < header_len + sizeof(struct tcphdr)) {
            printf("ERROR: Incomplete TCP packet\n");
            free(packet_data);
            return;
        }
        struct tcphdr *tcp_header = (struct tcphdr*)(packet_data + header_len);
        src_port = ntohs(tcp_header->source);
        dst_port = ntohs(tcp_header->dest);
        
        size_t tcp_header_len = tcp_header->doff * 4;
        if (bin_len > header_len + tcp_header_len) {
            payload = packet_data + header_len + tcp_header_len;
            payload_len = bin_len - header_len - tcp_header_len;
        }
    } else if (protocol_type == IPPROTO_UDP) {
        if (bin_len < header_len + sizeof(struct udphdr)) {
            printf("ERROR: Incomplete UDP packet\n");
            free(packet_data);
            return;
        }
        struct udphdr *udp_header = (struct udphdr*)(packet_data + header_len);
        src_port = ntohs(udp_header->source);
        dst_port = ntohs(udp_header->dest);
        
        if (bin_len > header_len + sizeof(struct udphdr)) {
            payload = packet_data + header_len + sizeof(struct udphdr);
            payload_len = bin_len - header_len - sizeof(struct udphdr);
        }
    }
    
    // Manual protocol detection on payload
    char *detected_protocol = "Unknown";
    char *detected_category = "Unspecified";
    
    if (payload && payload_len > 0) {
        // BitTorrent detection - look for protocol handshake
        if (payload_len >= 20 && payload[0] == 19 && 
            memcmp(payload + 1, "BitTorrent protocol", 19) == 0) {
            detected_protocol = "BitTorrent";
            detected_category = "P2P";
        }
        // HTTP detection
        else if (payload_len >= 4 && (memcmp(payload, "GET ", 4) == 0 || 
                                      memcmp(payload, "POST", 4) == 0 ||
                                      memcmp(payload, "HTTP", 4) == 0)) {
            detected_protocol = "HTTP";
            detected_category = "Web";
        }
        // TLS detection - look for TLS handshake
        else if (payload_len >= 6 && payload[0] == 0x16 && payload[1] == 0x03) {
            detected_protocol = "TLS";
            detected_category = "Web";
        }
        // SSH detection
        else if (payload_len >= 4 && memcmp(payload, "SSH-", 4) == 0) {
            detected_protocol = "SSH";
            detected_category = "RemoteAccess";
        }
    }
    
    // Output results
    printf("{\n");
    printf("  \"src_ip\": \"%u.%u.%u.%u\",\n", 
           (src_ip >> 24) & 0xFF, (src_ip >> 16) & 0xFF,
           (src_ip >> 8) & 0xFF, src_ip & 0xFF);
    printf("  \"dst_ip\": \"%u.%u.%u.%u\",\n",
           (dst_ip >> 24) & 0xFF, (dst_ip >> 16) & 0xFF,
           (dst_ip >> 8) & 0xFF, dst_ip & 0xFF);
    printf("  \"src_port\": %u,\n", src_port);
    printf("  \"dst_port\": %u,\n", dst_port);
    printf("  \"protocol\": \"%s\",\n", detected_protocol);
    printf("  \"category\": \"%s\",\n", detected_category);
    printf("  \"payload_len\": %zu\n", payload_len);
    printf("}\n");
    
    free(packet_data);
}

int main(int argc, char *argv[]) {
    if (argc != 2) {
        printf("Usage: %s <hex_packet_data>\n", argv[0]);
        printf("Example: %s \"4500003c...\"\n", argv[0]);
        return 1;
    }
    
    init_ndpi();
    analyze_packet_from_hex(argv[1]);
    cleanup_ndpi();
    
    return 0;
}