568 lines
21 KiB
Python
568 lines
21 KiB
Python
"""
|
||
Main GUI window for StreamLens
|
||
"""
|
||
|
||
import sys
|
||
import os
|
||
from typing import Optional, List, TYPE_CHECKING
|
||
|
||
try:
|
||
from PySide6.QtWidgets import (
|
||
QMainWindow, QWidget, QVBoxLayout, QHBoxLayout, QSplitter,
|
||
QTableWidget, QTableWidgetItem, QTextEdit, QMenuBar, QMenu,
|
||
QFileDialog, QMessageBox, QProgressBar, QStatusBar, QLabel,
|
||
QHeaderView, QPushButton, QGroupBox
|
||
)
|
||
from PySide6.QtCore import Qt, QThread, Signal, QTimer
|
||
from PySide6.QtGui import QAction, QIcon, QFont
|
||
|
||
# Matplotlib integration - lazy loaded
|
||
matplotlib = None
|
||
FigureCanvas = None
|
||
NavigationToolbar = None
|
||
Figure = None
|
||
plt = None
|
||
|
||
def _ensure_matplotlib_gui_loaded():
|
||
"""Lazy load matplotlib for GUI mode"""
|
||
global matplotlib, FigureCanvas, NavigationToolbar, Figure, plt
|
||
|
||
if matplotlib is not None:
|
||
return True
|
||
|
||
try:
|
||
import matplotlib as mpl
|
||
matplotlib = mpl
|
||
matplotlib.use('Qt5Agg') # Use Qt backend for matplotlib
|
||
|
||
from matplotlib.backends.backend_qt5agg import FigureCanvasQTAgg as FC
|
||
from matplotlib.backends.backend_qt5agg import NavigationToolbar2QT as NT
|
||
from matplotlib.figure import Figure as Fig
|
||
import matplotlib.pyplot as pyplot
|
||
|
||
FigureCanvas = FC
|
||
NavigationToolbar = NT
|
||
Figure = Fig
|
||
plt = pyplot
|
||
|
||
return True
|
||
except ImportError as e:
|
||
print(f"Matplotlib GUI integration not available: {e}")
|
||
return False
|
||
|
||
except ImportError as e:
|
||
print(f"GUI dependencies not available: {e}")
|
||
sys.exit(1)
|
||
|
||
if TYPE_CHECKING:
|
||
from ..analysis.core import EthernetAnalyzer
|
||
from ..models.flow_stats import FlowStats
|
||
|
||
|
||
class PlotWidget(QWidget):
|
||
"""Widget containing matplotlib plot with toolbar"""
|
||
|
||
def __init__(self, parent=None):
|
||
super().__init__(parent)
|
||
self.figure = None
|
||
self.canvas = None
|
||
self.toolbar = None
|
||
self._initialized = False
|
||
|
||
# Just create basic layout initially
|
||
self.layout = QVBoxLayout()
|
||
self.setLayout(self.layout)
|
||
|
||
def _ensure_initialized(self):
|
||
"""Initialize matplotlib components when first needed"""
|
||
if self._initialized:
|
||
return True
|
||
|
||
if not _ensure_matplotlib_gui_loaded():
|
||
return False
|
||
|
||
self.figure = Figure(figsize=(12, 8))
|
||
self.canvas = FigureCanvas(self.figure)
|
||
self.toolbar = NavigationToolbar(self.canvas, self)
|
||
|
||
self.layout.addWidget(self.toolbar)
|
||
self.layout.addWidget(self.canvas)
|
||
|
||
self._initialized = True
|
||
self.clear_plot()
|
||
return True
|
||
|
||
def clear_plot(self):
|
||
"""Clear the plot and show initial message"""
|
||
if not self._ensure_initialized():
|
||
return
|
||
self.figure.clear()
|
||
|
||
# Add a centered message when no plot is shown
|
||
ax = self.figure.add_subplot(111)
|
||
ax.text(0.5, 0.5, 'Select a flow with Chapter 10 data to view signal visualization',
|
||
ha='center', va='center', transform=ax.transAxes, fontsize=12,
|
||
color='gray', style='italic')
|
||
ax.set_xlim(0, 1)
|
||
ax.set_ylim(0, 1)
|
||
ax.axis('off')
|
||
|
||
self.canvas.draw()
|
||
|
||
def plot_flow_signals(self, flow: 'FlowStats', signal_data_list: List, flow_key: str):
|
||
"""Plot Chapter 10 signals for a flow"""
|
||
if not self._ensure_initialized():
|
||
return
|
||
self.figure.clear()
|
||
|
||
if not signal_data_list:
|
||
# Show message when no data
|
||
ax = self.figure.add_subplot(111)
|
||
ax.text(0.5, 0.5, 'No Chapter 10 signal data found in selected flow',
|
||
ha='center', va='center', transform=ax.transAxes, fontsize=14)
|
||
ax.set_xlim(0, 1)
|
||
ax.set_ylim(0, 1)
|
||
ax.axis('off')
|
||
self.canvas.draw()
|
||
return
|
||
|
||
# Create subplots for each signal
|
||
n_signals = len(signal_data_list)
|
||
axes = []
|
||
|
||
for idx, signal_data in enumerate(signal_data_list):
|
||
ax = self.figure.add_subplot(n_signals, 1, idx + 1)
|
||
axes.append(ax)
|
||
|
||
# Plot each channel in the signal data
|
||
for channel_name, data in signal_data.channels.items():
|
||
ax.plot(signal_data.timestamps, data, label=channel_name, linewidth=0.8)
|
||
|
||
ax.set_xlabel('Time (s)')
|
||
ax.set_ylabel('Amplitude')
|
||
ax.grid(True, alpha=0.3)
|
||
ax.legend()
|
||
|
||
# Add metadata info
|
||
if signal_data.metadata and signal_data.metadata.channel_configs:
|
||
config_info = []
|
||
for ch_id, config in signal_data.metadata.channel_configs.items():
|
||
if 'units' in config:
|
||
config_info.append(f"CH{ch_id}: {config.get('units', 'Unknown')}")
|
||
if config_info:
|
||
ax.set_title(f"Channels: {', '.join(config_info)}")
|
||
|
||
self.figure.suptitle(f'Chapter 10 Signals - Flow: {flow_key}', fontsize=14)
|
||
self.figure.tight_layout()
|
||
self.canvas.draw()
|
||
|
||
|
||
class PCAPLoadThread(QThread):
|
||
"""Thread for loading PCAP files without blocking UI"""
|
||
|
||
progress_updated = Signal(int)
|
||
loading_finished = Signal(object) # Emits the analyzer
|
||
error_occurred = Signal(str)
|
||
|
||
def __init__(self, file_path: str, parent=None):
|
||
super().__init__(parent)
|
||
self.file_path = file_path
|
||
|
||
def run(self):
|
||
try:
|
||
from ..analysis.core import EthernetAnalyzer
|
||
from ..utils.pcap_loader import PCAPLoader
|
||
|
||
# Create analyzer
|
||
analyzer = EthernetAnalyzer()
|
||
|
||
# Load PCAP
|
||
loader = PCAPLoader(self.file_path)
|
||
if not loader.validate_file():
|
||
self.error_occurred.emit(f"Invalid PCAP file: {self.file_path}")
|
||
return
|
||
|
||
packets = loader.load_all()
|
||
analyzer.all_packets = packets
|
||
|
||
# Process packets with progress updates
|
||
total_packets = len(packets)
|
||
for i, packet in enumerate(packets, 1):
|
||
analyzer._process_single_packet(packet, i)
|
||
|
||
# Update progress every 1000 packets
|
||
if i % 1000 == 0 or i == total_packets:
|
||
progress = int((i / total_packets) * 100)
|
||
self.progress_updated.emit(progress)
|
||
|
||
# Calculate statistics
|
||
analyzer.calculate_statistics()
|
||
|
||
self.loading_finished.emit(analyzer)
|
||
|
||
except Exception as e:
|
||
self.error_occurred.emit(str(e))
|
||
|
||
|
||
class StreamLensMainWindow(QMainWindow):
|
||
"""Main application window"""
|
||
|
||
def __init__(self):
|
||
super().__init__()
|
||
self.analyzer: Optional['EthernetAnalyzer'] = None
|
||
self.current_file = None
|
||
self.loading_thread = None
|
||
|
||
self.setWindowTitle("StreamLens - Ethernet Traffic Analyzer")
|
||
self.setGeometry(100, 100, 1400, 900)
|
||
|
||
self.setup_ui()
|
||
self.setup_menus()
|
||
self.setup_status_bar()
|
||
|
||
def setup_ui(self):
|
||
"""Set up the main UI layout"""
|
||
central_widget = QWidget()
|
||
self.setCentralWidget(central_widget)
|
||
|
||
# Create main horizontal splitter
|
||
main_splitter = QSplitter(Qt.Horizontal)
|
||
|
||
# Left panel - Flow list and controls
|
||
left_panel = self.create_left_panel()
|
||
left_panel.setMaximumWidth(600)
|
||
left_panel.setMinimumWidth(400)
|
||
|
||
# Right panel - Plot area
|
||
right_panel = self.create_right_panel()
|
||
|
||
main_splitter.addWidget(left_panel)
|
||
main_splitter.addWidget(right_panel)
|
||
main_splitter.setStretchFactor(0, 0) # Don't stretch left panel
|
||
main_splitter.setStretchFactor(1, 1) # Stretch right panel
|
||
|
||
layout = QHBoxLayout()
|
||
layout.addWidget(main_splitter)
|
||
central_widget.setLayout(layout)
|
||
|
||
def create_left_panel(self):
|
||
"""Create the left panel with flow list and info"""
|
||
widget = QWidget()
|
||
layout = QVBoxLayout()
|
||
|
||
# File info group
|
||
file_group = QGroupBox("File Information")
|
||
file_layout = QVBoxLayout()
|
||
self.file_info_label = QLabel("No file loaded")
|
||
self.file_info_label.setWordWrap(True)
|
||
file_layout.addWidget(self.file_info_label)
|
||
file_group.setLayout(file_layout)
|
||
layout.addWidget(file_group)
|
||
|
||
# Flow list group
|
||
flows_group = QGroupBox("IP Flows (sorted by max sigma deviation)")
|
||
flows_layout = QVBoxLayout()
|
||
|
||
self.flows_table = QTableWidget()
|
||
self.flows_table.setColumnCount(5)
|
||
self.flows_table.setHorizontalHeaderLabels([
|
||
"Source → Destination", "Packets", "Max σ", "Protocols", "Frame Types"
|
||
])
|
||
|
||
# Configure table
|
||
header = self.flows_table.horizontalHeader()
|
||
header.setStretchLastSection(True)
|
||
header.setSectionResizeMode(0, QHeaderView.Stretch)
|
||
|
||
self.flows_table.setSelectionBehavior(QTableWidget.SelectRows)
|
||
self.flows_table.setAlternatingRowColors(True)
|
||
self.flows_table.itemSelectionChanged.connect(self.on_flow_selected)
|
||
|
||
flows_layout.addWidget(self.flows_table)
|
||
|
||
# Status label for visualization feedback
|
||
self.viz_status_label = QLabel("Select a flow to view signal visualization")
|
||
self.viz_status_label.setStyleSheet("color: gray; font-style: italic;")
|
||
flows_layout.addWidget(self.viz_status_label)
|
||
|
||
flows_group.setLayout(flows_layout)
|
||
layout.addWidget(flows_group)
|
||
|
||
widget.setLayout(layout)
|
||
return widget
|
||
|
||
def create_right_panel(self):
|
||
"""Create the right panel with plot area"""
|
||
widget = QWidget()
|
||
layout = QVBoxLayout()
|
||
|
||
plot_group = QGroupBox("Signal Visualization")
|
||
plot_layout = QVBoxLayout()
|
||
|
||
self.plot_widget = PlotWidget()
|
||
plot_layout.addWidget(self.plot_widget)
|
||
|
||
plot_group.setLayout(plot_layout)
|
||
layout.addWidget(plot_group)
|
||
|
||
widget.setLayout(layout)
|
||
return widget
|
||
|
||
def setup_menus(self):
|
||
"""Set up application menus"""
|
||
menubar = self.menuBar()
|
||
|
||
# File menu
|
||
file_menu = menubar.addMenu("File")
|
||
|
||
open_action = QAction("Open PCAP...", self)
|
||
open_action.setShortcut("Ctrl+O")
|
||
open_action.triggered.connect(self.open_pcap_file)
|
||
file_menu.addAction(open_action)
|
||
|
||
file_menu.addSeparator()
|
||
|
||
exit_action = QAction("Exit", self)
|
||
exit_action.setShortcut("Ctrl+Q")
|
||
exit_action.triggered.connect(self.close)
|
||
file_menu.addAction(exit_action)
|
||
|
||
# View menu
|
||
view_menu = menubar.addMenu("View")
|
||
|
||
refresh_action = QAction("Refresh", self)
|
||
refresh_action.setShortcut("F5")
|
||
refresh_action.triggered.connect(self.refresh_data)
|
||
view_menu.addAction(refresh_action)
|
||
|
||
def setup_status_bar(self):
|
||
"""Set up status bar"""
|
||
self.status_bar = QStatusBar()
|
||
self.setStatusBar(self.status_bar)
|
||
|
||
self.progress_bar = QProgressBar()
|
||
self.progress_bar.setVisible(False)
|
||
self.status_bar.addPermanentWidget(self.progress_bar)
|
||
|
||
self.status_bar.showMessage("Ready")
|
||
|
||
def open_pcap_file(self):
|
||
"""Open PCAP file dialog"""
|
||
file_path, _ = QFileDialog.getOpenFileName(
|
||
self,
|
||
"Open PCAP File",
|
||
"",
|
||
"PCAP Files (*.pcap *.pcapng);;All Files (*)"
|
||
)
|
||
|
||
if file_path:
|
||
self.load_pcap_file(file_path)
|
||
|
||
def load_pcap_file(self, file_path: str):
|
||
"""Load PCAP file in background thread"""
|
||
if self.loading_thread and self.loading_thread.isRunning():
|
||
return
|
||
|
||
self.current_file = file_path
|
||
self.status_bar.showMessage(f"Loading {os.path.basename(file_path)}...")
|
||
self.progress_bar.setVisible(True)
|
||
self.progress_bar.setValue(0)
|
||
|
||
# Disable UI during loading
|
||
self.flows_table.setEnabled(False)
|
||
|
||
# Start loading thread
|
||
self.loading_thread = PCAPLoadThread(file_path)
|
||
self.loading_thread.progress_updated.connect(self.progress_bar.setValue)
|
||
self.loading_thread.loading_finished.connect(self.on_pcap_loaded)
|
||
self.loading_thread.error_occurred.connect(self.on_loading_error)
|
||
self.loading_thread.start()
|
||
|
||
def on_pcap_loaded(self, analyzer: 'EthernetAnalyzer'):
|
||
"""Handle PCAP loading completion"""
|
||
self.analyzer = analyzer
|
||
self.progress_bar.setVisible(False)
|
||
self.flows_table.setEnabled(True)
|
||
|
||
# Update file info
|
||
summary = analyzer.get_summary()
|
||
file_info = f"File: {os.path.basename(self.current_file)}\n"
|
||
file_info += f"Packets: {summary['total_packets']:,}\n"
|
||
file_info += f"Flows: {summary['unique_flows']}\n"
|
||
file_info += f"IPs: {summary['unique_ips']}"
|
||
self.file_info_label.setText(file_info)
|
||
|
||
# Populate flows table
|
||
self.populate_flows_table()
|
||
|
||
self.status_bar.showMessage(f"Loaded {summary['total_packets']:,} packets from {os.path.basename(self.current_file)}")
|
||
|
||
def on_loading_error(self, error_message: str):
|
||
"""Handle loading error"""
|
||
self.progress_bar.setVisible(False)
|
||
self.flows_table.setEnabled(True)
|
||
|
||
QMessageBox.critical(self, "Loading Error", f"Error loading PCAP file:\n{error_message}")
|
||
self.status_bar.showMessage("Error loading file")
|
||
|
||
def populate_flows_table(self):
|
||
"""Populate the flows table with data"""
|
||
if not self.analyzer:
|
||
return
|
||
|
||
summary = self.analyzer.get_summary()
|
||
flows_list = list(summary['flows'].values())
|
||
|
||
# Sort by maximum sigma deviation
|
||
flows_list.sort(key=lambda x: (
|
||
self.analyzer.statistics_engine.get_max_sigma_deviation(x),
|
||
x.frame_count
|
||
), reverse=True)
|
||
|
||
self.flows_table.setRowCount(len(flows_list))
|
||
|
||
for row, flow in enumerate(flows_list):
|
||
# Source -> Destination
|
||
flow_item = QTableWidgetItem(f"{flow.src_ip} → {flow.dst_ip}")
|
||
self.flows_table.setItem(row, 0, flow_item)
|
||
|
||
# Packets
|
||
packets_item = QTableWidgetItem(str(flow.frame_count))
|
||
packets_item.setData(Qt.UserRole, flow) # Store flow object
|
||
self.flows_table.setItem(row, 1, packets_item)
|
||
|
||
# Max sigma deviation
|
||
max_sigma = self.analyzer.statistics_engine.get_max_sigma_deviation(flow)
|
||
sigma_item = QTableWidgetItem(f"{max_sigma:.2f}σ")
|
||
self.flows_table.setItem(row, 2, sigma_item)
|
||
|
||
# Protocols
|
||
protocols = ", ".join(flow.protocols)
|
||
if flow.detected_protocol_types:
|
||
protocols += f" ({', '.join(flow.detected_protocol_types)})"
|
||
protocols_item = QTableWidgetItem(protocols)
|
||
self.flows_table.setItem(row, 3, protocols_item)
|
||
|
||
# Frame types
|
||
frame_types = ", ".join(flow.frame_types.keys())
|
||
frame_types_item = QTableWidgetItem(frame_types)
|
||
self.flows_table.setItem(row, 4, frame_types_item)
|
||
|
||
# Resize columns to content
|
||
self.flows_table.resizeColumnsToContents()
|
||
|
||
def on_flow_selected(self):
|
||
"""Handle flow selection and automatically render plots"""
|
||
selected_rows = self.flows_table.selectionModel().selectedRows()
|
||
|
||
if not selected_rows:
|
||
# No selection - clear plot and show message
|
||
self.plot_widget.clear_plot()
|
||
self.viz_status_label.setText("Select a flow to view signal visualization")
|
||
self.viz_status_label.setStyleSheet("color: gray; font-style: italic;")
|
||
return
|
||
|
||
# Get selected flow and automatically visualize
|
||
self.visualize_selected_flow()
|
||
|
||
def visualize_selected_flow(self):
|
||
"""Visualize the selected flow"""
|
||
if not self.analyzer:
|
||
return
|
||
|
||
selected_rows = self.flows_table.selectionModel().selectedRows()
|
||
if not selected_rows:
|
||
return
|
||
|
||
# Get selected flow
|
||
row = selected_rows[0].row()
|
||
flow_item = self.flows_table.item(row, 1) # Packets column has flow data
|
||
flow = flow_item.data(Qt.UserRole)
|
||
|
||
if not flow:
|
||
return
|
||
|
||
# Check if flow has Chapter 10 data
|
||
has_ch10 = any('CH10' in ft or 'TMATS' in ft for ft in flow.frame_types.keys())
|
||
|
||
if not has_ch10:
|
||
# Clear plot and show informative message
|
||
self.plot_widget.clear_plot()
|
||
self.viz_status_label.setText("Selected flow does not contain Chapter 10 telemetry data")
|
||
self.viz_status_label.setStyleSheet("color: orange; font-style: italic;")
|
||
return
|
||
|
||
# Get flow packets
|
||
flow_packets = self.get_flow_packets(flow)
|
||
|
||
if not flow_packets:
|
||
# Clear plot and show error message
|
||
self.plot_widget.clear_plot()
|
||
self.viz_status_label.setText("No packets found for selected flow")
|
||
self.viz_status_label.setStyleSheet("color: red; font-style: italic;")
|
||
return
|
||
|
||
# Use signal visualizer to extract and process signals
|
||
from ..utils.signal_visualizer import signal_visualizer
|
||
|
||
flow_key = f"{flow.src_ip} → {flow.dst_ip}"
|
||
self.status_bar.showMessage(f"Extracting signals from flow {flow_key}...")
|
||
self.viz_status_label.setText("Processing Chapter 10 telemetry data...")
|
||
self.viz_status_label.setStyleSheet("color: blue; font-style: italic;")
|
||
|
||
try:
|
||
# Extract signals
|
||
tmats_metadata = signal_visualizer._extract_tmats_from_flow(flow_packets)
|
||
signal_data_list = signal_visualizer._extract_signals_from_flow(flow_packets, tmats_metadata)
|
||
|
||
# Plot in GUI
|
||
self.plot_widget.plot_flow_signals(flow, signal_data_list, flow_key)
|
||
|
||
if signal_data_list:
|
||
self.status_bar.showMessage(f"Plotted {len(signal_data_list)} signal(s) from {flow_key}")
|
||
self.viz_status_label.setText(f"Displaying {len(signal_data_list)} signal channel(s)")
|
||
self.viz_status_label.setStyleSheet("color: green; font-weight: bold;")
|
||
else:
|
||
self.status_bar.showMessage(f"No decodable signal data found in {flow_key}")
|
||
self.viz_status_label.setText("No decodable signal data found in selected flow")
|
||
self.viz_status_label.setStyleSheet("color: orange; font-style: italic;")
|
||
|
||
except Exception as e:
|
||
# Show error in status label instead of popup
|
||
self.plot_widget.clear_plot()
|
||
self.viz_status_label.setText(f"Visualization error: {str(e)}")
|
||
self.viz_status_label.setStyleSheet("color: red; font-style: italic;")
|
||
self.status_bar.showMessage("Visualization error")
|
||
|
||
def get_flow_packets(self, flow: 'FlowStats') -> List:
|
||
"""Get all packets for a specific flow"""
|
||
if not self.analyzer or not self.analyzer.all_packets:
|
||
return []
|
||
|
||
flow_packets = []
|
||
|
||
for packet in self.analyzer.all_packets:
|
||
try:
|
||
if hasattr(packet, 'haslayer'):
|
||
from scapy.all import IP
|
||
if packet.haslayer(IP):
|
||
ip_layer = packet[IP]
|
||
if ip_layer.src == flow.src_ip and ip_layer.dst == flow.dst_ip:
|
||
flow_packets.append(packet)
|
||
except:
|
||
continue
|
||
|
||
return flow_packets
|
||
|
||
def refresh_data(self):
|
||
"""Refresh the current data"""
|
||
if self.current_file:
|
||
self.load_pcap_file(self.current_file)
|
||
|
||
def closeEvent(self, event):
|
||
"""Handle application close"""
|
||
if self.loading_thread and self.loading_thread.isRunning():
|
||
self.loading_thread.quit()
|
||
self.loading_thread.wait()
|
||
|
||
event.accept() |