Files
StreamLens/analyzer/utils/pcap_loader.py

96 lines
3.2 KiB
Python

"""
PCAP file loading utilities
"""
import sys
from typing import List, Iterator, Optional
try:
from scapy.all import rdpcap, PcapReader, Packet
except ImportError:
print("Error: scapy library required. Install with: pip install scapy")
sys.exit(1)
class PCAPLoader:
"""Utility class for loading PCAP files"""
def __init__(self, file_path: str):
self.file_path = file_path
self._packet_count: Optional[int] = None
def load_all(self) -> List[Packet]:
"""Load all packets from the PCAP file"""
try:
packets = rdpcap(self.file_path)
self._packet_count = len(packets)
return packets
except Exception as e:
raise IOError(f"Error loading PCAP file {self.file_path}: {e}")
def load_streaming(self, chunk_size: int = 1000) -> Iterator[List[Packet]]:
"""Load packets in chunks for memory efficiency"""
try:
with PcapReader(self.file_path) as pcap_reader:
chunk = []
for packet in pcap_reader:
chunk.append(packet)
if len(chunk) >= chunk_size:
yield chunk
chunk = []
# Yield remaining packets
if chunk:
yield chunk
except Exception as e:
raise IOError(f"Error streaming PCAP file {self.file_path}: {e}")
def get_packet_count(self) -> Optional[int]:
"""Get the total number of packets (if loaded)"""
return self._packet_count
def validate_file(self) -> bool:
"""Validate that the file exists and is a valid PCAP"""
try:
with PcapReader(self.file_path) as pcap_reader:
# Try to read first packet
next(iter(pcap_reader))
return True
except (IOError, StopIteration):
return False
except Exception:
return False
@staticmethod
def get_file_info(file_path: str) -> dict:
"""Get basic information about a PCAP file"""
try:
packet_count = 0
first_timestamp = None
last_timestamp = None
total_bytes = 0
with PcapReader(file_path) as pcap_reader:
for packet in pcap_reader:
packet_count += 1
total_bytes += len(packet)
if first_timestamp is None:
first_timestamp = float(packet.time)
last_timestamp = float(packet.time)
duration = (last_timestamp - first_timestamp) if first_timestamp and last_timestamp else 0
return {
'file_path': file_path,
'packet_count': packet_count,
'total_bytes': total_bytes,
'duration_seconds': duration,
'first_timestamp': first_timestamp,
'last_timestamp': last_timestamp,
'avg_packet_rate': packet_count / duration if duration > 0 else 0
}
except Exception as e:
return {'error': str(e)}