129 lines
4.1 KiB
Python
129 lines
4.1 KiB
Python
"""
|
|
Live network capture utilities
|
|
"""
|
|
|
|
import sys
|
|
import threading
|
|
import time
|
|
from typing import Callable, Optional, List
|
|
|
|
try:
|
|
from scapy.all import sniff, Packet, get_if_list
|
|
except ImportError:
|
|
print("Error: scapy library required. Install with: pip install scapy")
|
|
sys.exit(1)
|
|
|
|
|
|
class LiveCapture:
|
|
"""Utility class for live network capture"""
|
|
|
|
def __init__(self, interface: Optional[str] = None, filter_str: Optional[str] = None):
|
|
self.interface = interface
|
|
self.filter_str = filter_str
|
|
self.is_capturing = False
|
|
self.stop_requested = False
|
|
self.packet_count = 0
|
|
self.capture_thread: Optional[threading.Thread] = None
|
|
self.packet_handlers: List[Callable[[Packet, int], None]] = []
|
|
|
|
def add_packet_handler(self, handler: Callable[[Packet, int], None]) -> None:
|
|
"""Add a packet handler function"""
|
|
self.packet_handlers.append(handler)
|
|
|
|
def remove_packet_handler(self, handler: Callable[[Packet, int], None]) -> None:
|
|
"""Remove a packet handler function"""
|
|
if handler in self.packet_handlers:
|
|
self.packet_handlers.remove(handler)
|
|
|
|
def start_capture(self, threaded: bool = True) -> None:
|
|
"""Start packet capture"""
|
|
if self.is_capturing:
|
|
raise RuntimeError("Capture is already running")
|
|
|
|
self.stop_requested = False
|
|
self.packet_count = 0
|
|
|
|
if threaded:
|
|
self.capture_thread = threading.Thread(target=self._capture_loop, daemon=True)
|
|
self.capture_thread.start()
|
|
else:
|
|
self._capture_loop()
|
|
|
|
def stop_capture(self) -> None:
|
|
"""Stop packet capture"""
|
|
self.stop_requested = True
|
|
|
|
if self.capture_thread and self.capture_thread.is_alive():
|
|
self.capture_thread.join(timeout=5.0)
|
|
if self.capture_thread.is_alive():
|
|
print("Warning: Capture thread did not stop gracefully")
|
|
|
|
def _capture_loop(self) -> None:
|
|
"""Main capture loop"""
|
|
self.is_capturing = True
|
|
|
|
try:
|
|
def packet_handler(packet: Packet) -> None:
|
|
if self.stop_requested:
|
|
return
|
|
|
|
self.packet_count += 1
|
|
|
|
# Call all registered handlers
|
|
for handler in self.packet_handlers:
|
|
try:
|
|
handler(packet, self.packet_count)
|
|
except Exception as e:
|
|
print(f"Error in packet handler: {e}")
|
|
|
|
sniff(
|
|
iface=self.interface,
|
|
filter=self.filter_str,
|
|
prn=packet_handler,
|
|
stop_filter=lambda x: self.stop_requested
|
|
)
|
|
|
|
except Exception as e:
|
|
print(f"Error during live capture: {e}")
|
|
finally:
|
|
self.is_capturing = False
|
|
|
|
def get_capture_stats(self) -> dict:
|
|
"""Get capture statistics"""
|
|
return {
|
|
'is_capturing': self.is_capturing,
|
|
'packet_count': self.packet_count,
|
|
'interface': self.interface,
|
|
'filter': self.filter_str
|
|
}
|
|
|
|
@staticmethod
|
|
def list_interfaces() -> List[str]:
|
|
"""List available network interfaces"""
|
|
try:
|
|
return get_if_list()
|
|
except Exception as e:
|
|
print(f"Error listing interfaces: {e}")
|
|
return []
|
|
|
|
@staticmethod
|
|
def validate_interface(interface: str) -> bool:
|
|
"""Validate that an interface exists"""
|
|
try:
|
|
available_interfaces = LiveCapture.list_interfaces()
|
|
return interface in available_interfaces
|
|
except Exception:
|
|
return False
|
|
|
|
@staticmethod
|
|
def test_capture_permissions() -> bool:
|
|
"""Test if we have permissions for packet capture"""
|
|
try:
|
|
# Try a very short capture to test permissions
|
|
def dummy_handler(packet):
|
|
pass
|
|
|
|
sniff(count=1, timeout=1, prn=dummy_handler)
|
|
return True
|
|
except Exception:
|
|
return False |