packet_capture.py

#!/usr/bin/env python3
"""
DeepNet Packet Capture Engine
Handles packet sniffing and analysis
"""

import threading
import time
import queue
from datetime import datetime
from typing import Callable, Optional, List, Dict, Any
import json

try:
    from scapy.all import sniff, get_if_list, conf # get_if_list imported here
    from scapy.layers.inet import IP, TCP, UDP, ICMP
    from scapy.layers.inet6 import IPv6
    from scapy.layers.l2 import Ether, ARP
    from scapy.layers.dns import DNS
except ImportError:
    print("Error: Scapy not found. Install with: pip install scapy")
    exit(1)

from utils import format_bytes, get_protocol_name, extract_payload

class PacketCapture:
    """Main packet capture class"""
    
    def __init__(self, interface: Optional[str] = None, 
                 filter_str: Optional[str] = None,
                 callback: Optional[Callable] = None):
        self.interface = interface or self._get_default_interface()
        self.filter_str = filter_str
        self.callback = callback
        self.running = False
        self.packets = []
        self.packet_queue = queue.Queue()
        self.packet_count = 0
        self.capture_thread = None
        self.start_time = None
        self.statistics = {
            'total_packets': 0,
            'protocols': {},
            'data_transferred': 0,
            'unique_ips': set(),
            'packet_sizes': [],
            'most_active_src': {}, # Added for analysis
            'most_active_dst': {}, # Added for analysis
            'avg_packet_size': 0.0,
            'capture_duration': '0s'
        }
    
    def _get_default_interface(self) -> str:
        """Get the default network interface"""
        try:
            interfaces = get_if_list()
            # Filter out loopback and inactive interfaces
            active_interfaces = [iface for iface in interfaces 
                               if not iface.startswith('lo') and 
                               iface not in ['any', 'bluetooth-monitor']]
            
            if active_interfaces:
                return active_interfaces[0]
            else:
                return interfaces[0] if interfaces else 'any'
        except Exception:
            return 'any'
    
    def get_available_interfaces(self) -> List[str]:
        """Get list of available network interfaces"""
        try:
            return get_if_list()
        except Exception:
            return ['any']
    
    def start_capture(self, interface: str, filter_str: str, count: int = 0) -> None:
        """Start packet capture in a separate thread"""
        if self.running:
            return
        
        self.interface = interface # Use provided interface
        self.filter_str = filter_str # Use provided filter
        self.running = True
        self.start_time = datetime.now()
        self.packet_count = 0
        self.packets.clear()
        
        # Reset statistics
        self.statistics = {
            'total_packets': 0,
            'protocols': {},
            'data_transferred': 0,
            'unique_ips': set(),
            'packet_sizes': [],
            'most_active_src': {},
            'most_active_dst': {},
            'avg_packet_size': 0.0,
            'capture_duration': '0s'
        }
        
        self.capture_thread = threading.Thread(
            target=self._capture_packets,
            args=(count,),
            daemon=True
        )
        self.capture_thread.start()
    
    def stop_capture(self) -> None:
        """Stop packet capture"""
        self.running = False
        if self.capture_thread:
            self.capture_thread.join(timeout=2)
    
    def _capture_packets(self, count: int) -> None:
        """Internal method to capture packets"""
        try:
            sniff(
                iface=self.interface,
                filter=self.filter_str,
                count=count if count > 0 else 0,
                prn=self._process_packet,
                stop_filter=lambda x: not self.running,
                store=0 # Do not store packets in Scapy's internal list to save memory
            )
        except Exception as e:
            if self.callback:
                self.callback({
                    'type': 'error',
                    'message': f"Capture error: {str(e)}"
                })
    
    def _process_packet(self, packet) -> None:
        """Process each captured packet"""
        if not self.running:
            return
        
        try:
            packet_info = self._analyze_packet(packet)
            self.packet_queue.put(packet_info)
            
        except Exception as e:
            if self.callback:
                self.callback({
                    'type': 'error',
                    'message': f"Packet processing error: {str(e)}"
                })
    
    def process_queue(self):
        """Process packets from queue"""
        while not self.packet_queue.empty():
            packet_info = self.packet_queue.get()
            
            # Update statistics
            self._update_statistics(packet_info)
            
            # Store packet
            self.packets.append(packet_info)
            self.packet_count += 1
            
            # Call callback if provided, include statistics for GUI updates
            if self.callback:
                self.callback({
                    'type': 'packet',
                    'packet': packet_info,
                    'statistics': self.get_statistics()
                })
    
    def _analyze_packet(self, packet) -> Dict[str, Any]:
        """Analyze a single packet and extract information"""
        packet_info = {
            'timestamp': datetime.now().strftime('%Y-%m-%d %H:%M:%S.%f')[:-3],
            'size': len(packet),
            'protocol': 'Unknown',
            'src_ip': 'N/A',
            'dst_ip': 'N/A',
            'src_port': 'N/A',
            'dst_port': 'N/A',
            'payload': '',
            'flags': '',
            'summary': str(packet.summary()),
            'hex': bytes(packet).hex() # Add hex dump
        }
        
        # Ethernet layer
        if packet.haslayer(Ether):
            packet_info['src_mac'] = packet[Ether].src
            packet_info['dst_mac'] = packet[Ether].dst

        # IP layer
        if packet.haslayer(IP):
            ip_layer = packet[IP]
            packet_info['src_ip'] = ip_layer.src
            packet_info['dst_ip'] = ip_layer.dst
            packet_info['protocol'] = get_protocol_name(ip_layer.proto)
            packet_info['ttl'] = ip_layer.ttl
            packet_info['length'] = ip_layer.len

        # IPv6 layer
        elif packet.haslayer(IPv6):
            ipv6_layer = packet[IPv6]
            packet_info['src_ip'] = ipv6_layer.src
            packet_info['dst_ip'] = ipv6_layer.dst
            packet_info['protocol'] = 'IPv6'

        # TCP layer
        if packet.haslayer(TCP):
            tcp_layer = packet[TCP]
            packet_info['src_port'] = tcp_layer.sport
            packet_info['dst_port'] = tcp_layer.dport
            packet_info['protocol'] = 'TCP'
            packet_info['flags'] = self._get_tcp_flags(tcp_layer)
            packet_info['seq'] = tcp_layer.seq
            packet_info['ack'] = tcp_layer.ack
            packet_info['window'] = tcp_layer.window

        # UDP layer
        elif packet.haslayer(UDP):
            udp_layer = packet[UDP]
            packet_info['src_port'] = udp_layer.sport
            packet_info['dst_port'] = udp_layer.dport
            packet_info['protocol'] = 'UDP'
            packet_info['length'] = udp_layer.len

        # ICMP layer
        elif packet.haslayer(ICMP):
            packet_info['protocol'] = 'ICMP'
            packet_info['type'] = packet[ICMP].type
            packet_info['code'] = packet[ICMP].code

        # ARP layer
        elif packet.haslayer(ARP):
            arp_layer = packet[ARP]
            packet_info['protocol'] = 'ARP'
            packet_info['src_ip'] = arp_layer.psrc
            packet_info['dst_ip'] = arp_layer.pdst
            packet_info['hwsrc'] = arp_layer.hwsrc
            packet_info['hwdst'] = arp_layer.hwdst

        # DNS layer (often within UDP/TCP)
        if packet.haslayer(DNS):
            dns_layer = packet[DNS]
            packet_info['dns_query'] = dns_layer.qd.qname.decode() if dns_layer.qd else 'N/A'
            packet_info['protocol'] = 'DNS' # Override if DNS is present

        # Extract payload
        packet_info['payload'] = extract_payload(packet)

        return packet_info

    def _get_tcp_flags(self, tcp_layer) -> str:
        """Decode TCP flags"""
        flags = {
            'F': 'FIN', 'S': 'SYN', 'R': 'RST', 'P': 'PSH',
            'A': 'ACK', 'U': 'URG', 'E': 'ECE', 'C': 'CWR'
        }
        return ','.join([flags[f] for f in str(tcp_layer.flags)])

    def _update_statistics(self, packet_info: Dict[str, Any]) -> None:
        """Update capture statistics"""
        self.statistics['total_packets'] += 1
        self.statistics['data_transferred'] += packet_info['size']
        self.statistics['packet_sizes'].append(packet_info['size'])

        # Protocol distribution
        proto = packet_info['protocol']
        self.statistics['protocols'][proto] = self.statistics['protocols'].get(proto, 0) + 1

        # Unique IPs and Most Active IPs
        src_ip = packet_info['src_ip']
        dst_ip = packet_info['dst_ip']

        if src_ip != 'N/A':
            self.statistics['unique_ips'].add(src_ip)
            self.statistics['most_active_src'][src_ip] = self.statistics['most_active_src'].get(src_ip, 0) + 1
        
        if dst_ip != 'N/A':
            self.statistics['unique_ips'].add(dst_ip)
            self.statistics['most_active_dst'][dst_ip] = self.statistics['most_active_dst'].get(dst_ip, 0) + 1

        # Calculate average packet size
        if self.statistics['total_packets'] > 0:
            self.statistics['avg_packet_size'] = sum(self.statistics['packet_sizes']) / self.statistics['total_packets']
        
        # Calculate capture duration
        if self.start_time:
            duration = datetime.now() - self.start_time
            self.statistics['capture_duration'] = str(duration).split('.')[0] # Remove microseconds

    def get_statistics(self) -> Dict[str, Any]:
        """Get current capture statistics"""
        # Ensure unique_ips is a count, not the set itself, for reporting
        stats_copy = self.statistics.copy()
        stats_copy['unique_ips'] = len(self.statistics['unique_ips'])
        
        # Sort active IPs for display
        stats_copy['most_active_src'] = dict(sorted(self.statistics['most_active_src'].items(), 
                                                   key=lambda item: item[1], reverse=True))
        stats_copy['most_active_dst'] = dict(sorted(self.statistics['most_active_dst'].items(), 
                                                   key=lambda item: item[1], reverse=True))
        return stats_copy

    def get_packets(self) -> List[Dict[str, Any]]:
        """Get all captured packets"""
        return list(self.packets) # Return a copy to prevent external modification

    def clear_packets(self) -> None:
        """Clear all captured packets and reset statistics"""
        self.stop_capture() # Ensure capture is stopped before clearing
        self.packets.clear()
        self.packet_queue = queue.Queue() # Clear the queue as well
        self.packet_count = 0
        self.start_time = None
        self.statistics = {
            'total_packets': 0,
            'protocols': {},
            'data_transferred': 0,
            'unique_ips': set(),
            'packet_sizes': [],
            'most_active_src': {},
            'most_active_dst': {},
            'avg_packet_size': 0.0,
            'capture_duration': '0s'
        }

    def save_to_pcap(self, filename: str) -> None:
        """Save captured packets to PCAP file"""
        # This functionality requires storing raw Scapy packets, which is
        # currently disabled with store=0 for memory efficiency.
        # To enable PCAP saving, 'store=1' or similar handling would be needed
        # in the sniff function, along with converting stored dicts back to Scapy packets.
        # This is a placeholder for future implementation.
        raise NotImplementedError("Saving to PCAP is not currently supported in this version due to memory optimization.")
    
    def save_to_file(self, filename: str, format_type: str = 'json') -> None:
        """Save captured packets to file"""
        try:
            if format_type.lower() == 'json':
                with open(filename, 'w') as f:
                    json.dump(self.packets, f, indent=2)
            elif format_type.lower() == 'txt':
                with open(filename, 'w') as f:
                    for packet in self.packets:
                        f.write(f"Timestamp: {packet['timestamp']}\n")
                        f.write(f"Protocol: {packet['protocol']}\n")
                        f.write(f"Source: {packet['src_ip']}:{packet['src_port']}\n")
                        f.write(f"Destination: {packet['dst_ip']}:{packet['dst_port']}\n")
                        f.write(f"Size: {packet['size']} bytes\n")
                        f.write(f"Payload: {packet['payload'][:100]}...\n")
                        f.write("-" * 50 + "\n")
            else:
                raise ValueError("Unsupported format type. Choose 'json' or 'txt'.")
        except Exception as e:
            raise Exception(f"Error saving file: {str(e)}")