112 lines
4.0 KiB
Python
112 lines
4.0 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 with memory management"""
|
|
try:
|
|
# Check file size and warn for large files
|
|
import os
|
|
file_size = os.path.getsize(self.file_path)
|
|
if file_size > 100 * 1024 * 1024: # > 100MB
|
|
print(f"Warning: Large PCAP file ({file_size / (1024*1024):.1f} MB)")
|
|
print("Consider using streaming mode for better memory management")
|
|
|
|
# Force garbage collection before loading
|
|
import gc
|
|
gc.collect()
|
|
|
|
packets = rdpcap(self.file_path)
|
|
self._packet_count = len(packets)
|
|
|
|
print(f"Loaded {len(packets)} packets ({file_size / (1024*1024):.1f} MB)")
|
|
return packets
|
|
|
|
except MemoryError:
|
|
raise IOError(f"Out of memory loading PCAP file {self.file_path}. Try using streaming mode.")
|
|
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)} |