cli.py

#!/usr/bin/env python3
"""
DeepNet CLI Interface
Command line interface for network packet analysis
"""

import os
import sys
import time
import threading
from datetime import datetime
from typing import Optional, Dict, Any

from packet_capture import PacketCapture
from utils import (
    format_bytes, colorize_protocol, truncate_string, 
    format_timestamp, generate_summary_report, get_port_service
)

class DeepNetCLI:
    """Command Line Interface for DeepNet"""
    
    def __init__(self, interface: Optional[str] = None, count: int = 0,
                 filter_str: Optional[str] = None, output_file: Optional[str] = None,
                 verbose: bool = False):
        self.interface = interface
        self.count = count
        self.filter_str = filter_str
        self.output_file = output_file
        self.verbose = verbose
        self.capture = None
        self.running = False
        self.display_thread = None
        self.last_packet_count = 0
        self.process_thread = None  # Thread for processing packet queue
        
        # Terminal colors
        self.COLORS = {
            'GREEN': '\033[92m',
            'BLUE': '\033[94m',
            'YELLOW': '\033[93m',
            'RED': '\033[91m',
            'MAGENTA': '\033[95m',
            'CYAN': '\033[96m',
            'WHITE': '\033[97m',
            'BOLD': '\033[1m',
            'RESET': '\033[0m'
        }
    
    def run(self):
        """Run the CLI interface"""
        self.print_header()
        self.setup_capture()
        
        try:
            self.start_capture()
            self.run_interface()
        except KeyboardInterrupt:
            print(f"\n\n{self.COLORS['YELLOW']}[INFO] Stopping DeepNet...{self.COLORS['RESET']}")
        finally:
            self.cleanup()
    
    def print_header(self):
        """Print CLI header"""
        print(f"\n{self.COLORS['CYAN']}{'='*70}")
        print(f"{self.COLORS['BOLD']}DeepNet - Network Packet Analyzer (CLI Mode){self.COLORS['RESET']}")
        print(f"{self.COLORS['CYAN']}{'='*70}{self.COLORS['RESET']}")
        print(f"{self.COLORS['YELLOW']}⚠️  Educational Use Only - Monitor only networks you own{self.COLORS['RESET']}")
        print(f"{self.COLORS['CYAN']}{'='*70}{self.COLORS['RESET']}\n")
    
    def prompt_for_interface(self, interfaces):
        """Prompt the user to select an interface by number"""
        print(f"{self.COLORS['YELLOW']}Please select a network interface to use:{self.COLORS['RESET']}")
        for i, iface in enumerate(interfaces):
            print(f"  {i+1}. {iface}")
        while True:
            choice = input(f"Enter interface number [1]: ").strip()
            if choice == '':
                return interfaces[0]
            if choice.isdigit():
                idx = int(choice) - 1
                if 0 <= idx < len(interfaces):
                    return interfaces[idx]
            print(f"{self.COLORS['RED']}Invalid selection. Please enter a valid number.{self.COLORS['RESET']}")
    
    def setup_capture(self):
        """Setup packet capture"""
        print(f"{self.COLORS['BLUE']}[INFO] Initializing packet capture...{self.COLORS['RESET']}")
        
        self.capture = PacketCapture(
            interface=self.interface,
            filter_str=self.filter_str,
            callback=self.packet_callback
        )
        
        # Display configuration
        print(f"{self.COLORS['GREEN']}Configuration:{self.COLORS['RESET']}")
        print(f"  Interface: {self.capture.interface}")
        print(f"  Filter: {self.filter_str or 'None'}")
        print(f"  Count: {self.count if self.count > 0 else 'Unlimited'}")
        print(f"  Output: {self.output_file or 'None'}")
        print(f"  Verbose: {self.verbose}")
        
        # Show available interfaces
        interfaces = self.capture.get_available_interfaces() 
        print(f"\n{self.COLORS['BLUE']}Available interfaces:{self.COLORS['RESET']}")
        for i, iface in enumerate(interfaces):
            marker = " (selected)" if iface == self.capture.interface else ""
            print(f"  {i+1}. {iface}{marker}")
        print()
        
        # Prompt for interface if not provided
        if not self.interface:
            self.interface = self.prompt_for_interface(interfaces)
            self.capture.interface = self.interface
            print(f"{self.COLORS['GREEN']}Selected interface: {self.interface}{self.COLORS['RESET']}")
    
    def start_capture(self):
        """Start packet capture"""
        print(f"{self.COLORS['GREEN']}[INFO] Starting packet capture...{self.COLORS['RESET']}")
        print(f"{self.COLORS['YELLOW']}Press Ctrl+C to stop capture{self.COLORS['RESET']}\n")
        
        self.running = True
        # Ensure interface and filter_str are strings
        interface = self.interface if self.interface is not None else ''
        filter_str = self.filter_str if self.filter_str is not None else ''
        if self.capture:
            self.capture.start_capture(interface, filter_str, self.count)
        
        # Start background thread to process packet queue
        self.process_thread = threading.Thread(target=self.process_packets_loop, daemon=True)
        self.process_thread.start()
        
        # Start display thread for real-time updates
        if not self.verbose:
            self.display_thread = threading.Thread(target=self.display_summary, daemon=True)
            self.display_thread.start()
    
    def packet_callback(self, data: Dict[str, Any]):
        """Callback for packet events"""
        if data['type'] == 'packet':
            if self.verbose:
                self.display_packet(data['packet'])
        elif data['type'] == 'error':
            print(f"\n{self.COLORS['RED']}[ERROR] {data['message']}{self.COLORS['RESET']}")
    
    def display_packet(self, packet: Dict[str, Any]):
        """Display individual packet information"""
        timestamp = format_timestamp(packet['timestamp'])
        protocol = colorize_protocol(packet['protocol'])
        src = f"{packet['src_ip']}:{packet['src_port']}"
        dst = f"{packet['dst_ip']}:{packet['dst_port']}"
        size = packet['size']
        
        print(f"{timestamp} {protocol:15} {src:25} -> {dst:25} ({size:4} bytes)")
        
        # Show additional details in verbose mode
        if packet.get('flags'):
            print(f"  Flags: {packet['flags']}")
        if packet.get('payload') and len(packet['payload']) > 0:
            payload_preview = truncate_string(packet['payload'], 60)
            print(f"  Payload: {payload_preview}")
        print()
    
    def display_summary(self):
        """Display summary statistics in real-time"""
        while self.running: # Ensure loop terminates when capture stops
            try:
                # Clear screen and move cursor to top
                os.system('clear' if os.name == 'posix' else 'cls') 
                
                # Print header
                print(f"{self.COLORS['CYAN']}{'='*70}")
                print(f"{self.COLORS['BOLD']}DeepNet - Live Packet Capture{self.COLORS['RESET']}")
                print(f"{self.COLORS['CYAN']}{'='*70}{self.COLORS['RESET']}")
                
                # Get current statistics
                stats = self.capture.get_statistics() if self.capture else {
                    'total_packets': 0,
                    'data_transferred': 0,
                    'unique_ips': 0,
                    'capture_duration': '0s',
                    'avg_packet_size': 0.0,
                    'protocols': {},
                    'most_active_src': {},
                    'most_active_dst': {}
                }
                
                # Display statistics
                print(f"\n{self.COLORS['GREEN']}\U0001F4CA CAPTURE STATISTICS:{self.COLORS['RESET']}")
                print(f"  Total Packets: {stats['total_packets']}")
                print(f"  Data Transferred: {format_bytes(stats['data_transferred'])}")
                print(f"  Unique IPs: {stats['unique_ips']}")
                print(f"  Capture Duration: {stats['capture_duration']}")
                print(f"  Average Packet Size: {stats['avg_packet_size']:.2f} bytes")
                
                # Protocol distribution
                print(f"\n{self.COLORS['BLUE']}\U0001F4C8 PROTOCOL DISTRIBUTION:{self.COLORS['RESET']}")
                sorted_protocols = sorted(stats['protocols'].items(), 
                                          key=lambda item: item[1], reverse=True)
                total_packets = stats['total_packets'] if stats['total_packets'] > 0 else 1
                for proto, count in sorted_protocols:
                    percentage = (count / total_packets) * 100
                    print(f"  {proto:10}: {count:6} packets ({percentage:.1f}%)")
                
                # Active hosts (simplified for CLI)
                print(f"\n{self.COLORS['MAGENTA']}\U0001F310 ACTIVE HOSTS (TOP 5 SOURCE/DEST):{self.COLORS['RESET']}")
                print("  Top Source IPs:")
                for ip, count in list(stats['most_active_src'].items())[:5]:
                    print(f"    {ip:15}: {count:6} packets")
                print("  Top Destination IPs:")
                for ip, count in list(stats['most_active_dst'].items())[:5]:
                    print(f"    {ip:15}: {count:6} packets")

                # Packet rate
                current_packet_count = stats['total_packets']
                new_packets = current_packet_count - self.last_packet_count
                self.last_packet_count = current_packet_count
                
                print(f"\n{self.COLORS['YELLOW']}\u26A1 Live Packet Rate: {new_packets} packets/sec{self.COLORS['RESET']}")
                print(f"\n{self.COLORS['YELLOW']}Press Ctrl+C to stop capture{self.COLORS['RESET']}")

                time.sleep(1) # Update every 1 second
            except Exception as e:
                # Handle potential errors during display without crashing
                print(f"\n{self.COLORS['RED']}[ERROR] Display update failed: {str(e)}{self.COLORS['RESET']}")
                break # Exit display loop on error
    
    def process_packets_loop(self):
        """Continuously process packet queue while running"""
        while self.running:
            if self.capture:
                self.capture.process_queue()
            time.sleep(0.1)  # Process every 100ms
    
    def run_interface(self):
        """Keep the CLI interface running until stopped"""
        while self.running:
            time.sleep(0.5) # Keep main thread alive
    
    def stop_capture(self):
        """Stop the packet capture"""
        if self.running:
            print(f"\n{self.COLORS['YELLOW']}[INFO] Stopping capture...{self.COLORS['RESET']}")
            self.running = False
            if self.capture:
                self.capture.stop_capture()
            if self.display_thread and self.display_thread.is_alive():
                self.display_thread.join()
            if self.process_thread and self.process_thread.is_alive():
                self.process_thread.join()
            print(f"{self.COLORS['GREEN']}[INFO] Capture stopped.{self.COLORS['RESET']}")
    
    def cleanup(self):
        """Perform cleanup operations"""
        self.stop_capture()
        if self.output_file and self.capture: # Ensure capture object exists
            print(f"\n{self.COLORS['BLUE']}[INFO] Saving captured packets to {self.output_file}...{self.COLORS['RESET']}")
            self.save_packets()
        print(f"{self.COLORS['CYAN']}[INFO] DeepNet CLI Exited.{self.COLORS['RESET']}")

    def save_packets(self):
        """Save captured packets to a file"""
        if not self.capture or not self.capture.get_packets():
            print(f"{self.COLORS['YELLOW']}[WARN] No packets to save.{self.COLORS['RESET']}")
            return

        if not self.output_file:
            print(f"{self.COLORS['RED']}[ERROR] No output file specified. Use --output <filename> to save.{self.COLORS['RESET']}")
            return

        try:
            format_type = 'json'
            if self.output_file.lower().endswith('.txt'):
                format_type = 'txt'
            
            self.capture.save_to_file(self.output_file, format_type)
            
            packet_count = len(self.capture.get_packets())
            print(f"{self.COLORS['GREEN']}[SUCCESS] Saved {packet_count} packets to {self.output_file}{self.COLORS['RESET']}")
            
        except Exception as e:
            print(f"{self.COLORS['RED']}[ERROR] Failed to save packets: {str(e)}{self.COLORS['RESET']}")
    
    def print_help(self):
        """Print help information"""
        help_text = f"""
{self.COLORS['BOLD']}DeepNet CLI Commands:{self.COLORS['RESET']}

{self.COLORS['GREEN']}Keyboard Controls:{self.COLORS['RESET']}
  Ctrl+C    Stop packet capture
  
{self.COLORS['GREEN']}Filter Examples:{self.COLORS['RESET']}
  tcp                    Capture only TCP packets
  udp port 53           Capture DNS traffic
  host 192.168.1.1      Capture traffic to/from specific host
  tcp port 80 or 443    Capture HTTP/HTTPS traffic
  icmp                  Capture ICMP packets
  
{self.COLORS['GREEN']}Usage Examples:{self.COLORS['RESET']}
  python main.py cli --interface eth0 --count 100
  python main.py cli --filter "tcp port 80" --verbose
  python main.py cli --output capture.json --count 1000
  
{self.COLORS['YELLOW']}For more details, use: python main.py cli --help{self.COLORS['RESET']}
"""
        print(help_text)