import struct
import socket
import os
from collections import defaultdict

# Ethernet header length
ETH_HEADER_LENGTH = 14

# Constants for the PCAP Global Header format (assuming Ethernet, IPv4, and TCP/UDP)
PCAP_GLOBAL_HEADER_FORMAT = '@IHHiIII'

# PCAP Packet Header format
PCAP_PACKET_HEADER_FORMAT = '@IIII'


def ip_to_str(inet):
    """Convert binary IP address to a human-readable string."""
    try:
        return socket.inet_ntop(socket.AF_INET, inet)
    except ValueError:
        return False


def parse_pcap_file(pcap_file, target_ip):
    """Process and extract relevant packet data from the given pcap file."""
    with open(pcap_file, 'rb') as file:
        global_header = file.read(24)
        if len(global_header) < 24:
            raise ValueError('Invalid pcap file: global header length is less than 24 bytes.')

        packet_data_dict = {}

        while True:
            packet_header = file.read(16)
            if len(packet_header) < 16:
                break  # End of file

            timestamp_sec, timestamp_usec, included_len, original_len = struct.unpack(PCAP_PACKET_HEADER_FORMAT,
                                                                                      packet_header)
            packet_data = file.read(included_len)
            ethernet_header = packet_data[:ETH_HEADER_LENGTH]
            ethernet_fields = struct.unpack('!6s6sH', ethernet_header)
            ethernet_protocol = socket.ntohs(ethernet_fields[2])

            if ethernet_protocol == 8:  # IPv4
                ip_header = packet_data[ETH_HEADER_LENGTH:20 + ETH_HEADER_LENGTH]
                ip_fields = struct.unpack('!BBHHHBBH4s4s', ip_header)
                protocol_type = ip_fields[6]
                source_ip = socket.inet_ntoa(ip_fields[8])
                dest_ip = socket.inet_ntoa(ip_fields[9])

                if dest_ip == target_ip:
                    timestamp = timestamp_sec + (timestamp_usec / 1_000_000)
                    if protocol_type == 6:  # TCP
                        tcp_header = packet_data[
                                     ETH_HEADER_LENGTH + ip_fields[0] & 0xF * 4:ETH_HEADER_LENGTH + ip_fields[
                                         0] & 0xF * 4 + 20]
                        tcp_fields = struct.unpack('!HH', tcp_header[0:4])  # only need source and destination ports
                        source_port, destination_port = tcp_fields[0], tcp_fields[1]
                        key = (source_ip, destination_port)
                        packet_data_dict.setdefault(key, []).append((timestamp, source_ip))

                    elif protocol_type == 17:  # UDP
                        udp_header = packet_data[
                                     ETH_HEADER_LENGTH + ip_fields[0] & 0xF * 4:ETH_HEADER_LENGTH + ip_fields[
                                         0] & 0xF * 4 + 8]
                        udp_fields = struct.unpack('!HH', udp_header[0:4])
                        source_port, destination_port = udp_fields[0], udp_fields[1]
                        key = (source_ip, destination_port)
                        packet_data_dict.setdefault(key, []).append((timestamp, source_ip))

    return packet_data_dict


def identify_probing_packets(packet_data, probe_time_width, min_probes):
    """Identify probing packets based on the specified time width and minimum probe requirement."""
    probing_results = defaultdict(list)

    for (src_ip, dest_port), packets in packet_data.items():
        packets.sort()  # Sort by timestamp
        current_group = []

        for pkt in packets:
            if not current_group:
                current_group.append(pkt)
            else:
                # Check if the packets are within the probing time window
                if (pkt[0] - current_group[-1][0] <= probe_time_width):
                    current_group.append(pkt)
                else:
                    if len(current_group) >= min_probes:
                        probing_results[(src_ip, dest_port)].append(current_group)
                    current_group = [pkt]

        # Final check for the last group
        if len(current_group) >= min_probes:
            probing_results[(src_ip, dest_port)].append(current_group)

    return probing_results


def identify_scanning_packets(packet_data, scan_time_width, min_points_in_cluster):
    """Identify scanning packets based on time width and minimum points required in a cluster."""
    clusters_found = defaultdict(list)

    for (src_ip, dest_port), packets in packet_data.items():
        for packet in packets:
            # Consolidate scanning logic
            clusters_found[src_ip].append(packet)

    # Filter out clusters based on the minimum points requirement
    return {src_ip: clusters for src_ip, clusters in clusters_found.items() if len(clusters) >= min_points_in_cluster}


# Hardcoded path to the PCAP files folder
PCAP_FOLDER_PATH = r"C:\Users\tj471\PycharmProjects\Lab_Ass_4\Lab3-pcap-2"  # Update this path to your PCAP files location

# Get user input for parameters
target_ip = input("Enter the target IP address: ")
probe_width = int(input("Enter the width for probing in seconds: "))
min_probe_packets = int(input("Enter the minimum number of packets for probing: "))
scan_width = int(input("Enter the width for scanning in port numbers: "))
min_scan_packets = int(input("Enter the minimum number of packets for scanning: "))

# Load all PCAP files from the specified folder
pcap_files = [f for f in os.listdir(PCAP_FOLDER_PATH) if f.endswith('.pcap')]
all_probes_tcp = []
all_probes_udp = []
all_scans_tcp = []
all_scans_udp = []

# Process each PCAP file in the folder
for pcap_file in pcap_files:
    full_path = os.path.join(PCAP_FOLDER_PATH, pcap_file)
    print(f"Processing {full_path}...")

    # Process the pcap file and extract data
    packet_data = parse_pcap_file(full_path, target_ip)
    probes_detected = identify_probing_packets(packet_data, probe_width, min_probe_packets)
    scanned_clusters_detected = identify_scanning_packets(packet_data, scan_width, min_scan_packets)

    # Aggregate results for TCP and UDP separately
    for (src_ip, dest_port), probe_groups in probes_detected.items():
        for group in probe_groups:
            if all(pkt[1] == src_ip for pkt in group):  # Check if all source IPs are the same
                all_probes_tcp.append((src_ip, dest_port, len(group)))

    # Assume that scanning also can be TCP or UDP
    for src_ip, scan_group in scanned_clusters_detected.items():
        all_scans_tcp.append((src_ip, len(scan_group)))  # Use only source IP for scanning

# Output results to the console for probing and scanning
print(f"\nReports of Probing with TCP:")
for src_ip, dest_port, attempts in all_probes_tcp:
    print(f'Probing from {src_ip} to Port {dest_port} - Total Attempts: {attempts}')

print(f"\nReports of Probing with UDP:")
# Assuming UDP probe results were intended to be captured similarly
for (src_ip, dest_port), packets in packet_data.items():
    if packets:  # This is just to illustrate; you may need to refine the UDP capturing logic
        print(f'Probing from {src_ip} to Port {dest_port} - Total Attempts: {len(packets)}')

print(f"\nReports of Scanning with TCP:")
for src_ip, attempts in all_scans_tcp:
    print(f'Scanning from {src_ip} - Total Attempts: {attempts}')

print(f"\nReports of Scanning with UDP:")
# You can implement a similar logic for UDP scanning if applicable
for src_ip, attempts in all_scans_tcp:  # This is a placeholder. Adjust according to your actual UDP scan data
    print(f'Scanning from {src_ip} - Total Attempts: {attempts}')

print("\nAnalysis complete.")