Files
StreamLens/analyzer/utils/live_capture.py

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