diff --git a/c_extension/arpscanner.c b/c_extension/arpscanner.c index 563670a..4ddb006 100644 --- a/c_extension/arpscanner.c +++ b/c_extension/arpscanner.c @@ -51,7 +51,7 @@ typedef struct { arp_response_t *response; struct in_addr target_ip; int timeout_ms; - int finished; + volatile int finished; } capture_thread_args_t; /* @@ -144,7 +144,7 @@ static PyObject *perform_arp_scan(PyObject *self, PyObject *args) { // Open pcap handle char errbuf[PCAP_ERRBUF_SIZE]; - pcap_t *handle = pcap_open_live(iface, 65536, 1, timeout_ms, errbuf); + pcap_t *handle = pcap_open_live(iface, 65536, 1, 10, errbuf); if (!handle) { PyErr_SetString(PyExc_RuntimeError, errbuf); return NULL; @@ -154,7 +154,7 @@ static PyObject *perform_arp_scan(PyObject *self, PyObject *args) { struct bpf_program fp; char filter_exp[64]; snprintf(filter_exp, sizeof(filter_exp), "arp src host %s", dst_ip_str); - if (pcap_compile(handle, &fp, filter_exp, 0, PCAP_NETMASK_UNKNOWN) == -1) { + if (pcap_compile(handle, &fp, filter_exp, 1, PCAP_NETMASK_UNKNOWN) == -1) { pcap_close(handle); PyErr_SetString(PyExc_RuntimeError, "Failed to compile filter"); return NULL; @@ -207,18 +207,20 @@ static PyObject *perform_arp_scan(PyObject *self, PyObject *args) { return NULL; } + // Release the GIL for the blocking network operations so other Python + // threads (e.g. the ThreadPoolExecutor workers) can run concurrently. + int send_failed; + Py_BEGIN_ALLOW_THREADS + // Send ARP request - if (pcap_sendpacket(handle, packet, sizeof(packet)) != 0) { - pthread_cancel(tid); - pcap_close(handle); - PyErr_SetString(PyExc_RuntimeError, "Failed to send ARP packet"); - return NULL; - } + send_failed = (pcap_sendpacket(handle, packet, sizeof(packet)) != 0); - // Wait for response or timeout - while (!thread_args.finished && timeout_ms > 0) { - usleep(10000); // Sleep for 10ms - timeout_ms -= 10; + if (!send_failed) { + // Wait for response or timeout + while (!thread_args.finished && timeout_ms > 0) { + usleep(10000); // Sleep for 10ms + timeout_ms -= 10; + } } // Clean up @@ -226,6 +228,13 @@ static PyObject *perform_arp_scan(PyObject *self, PyObject *args) { pthread_join(tid, NULL); pcap_close(handle); + Py_END_ALLOW_THREADS + + if (send_failed) { + PyErr_SetString(PyExc_RuntimeError, "Failed to send ARP packet"); + return NULL; + } + // Return results if (response.found) { char mac_str[18]; diff --git a/core/arp_scanner.py b/core/arp_scanner.py index dc17400..f7385a8 100644 --- a/core/arp_scanner.py +++ b/core/arp_scanner.py @@ -11,18 +11,18 @@ import netifaces from PySide6.QtCore import Qt, QThread, Signal, Slot # pylint: disable=E0611 from PySide6.QtGui import QColor, QFont, QIcon # pylint: disable=E0611 -from PySide6.QtWidgets import (QDialog, QFileDialog, # pylint: disable=E0611 - QLabel, QLineEdit, QListWidget, QListWidgetItem, - QMainWindow, QMessageBox, QProgressBar, - QPushButton, QSplitter, QTextEdit, QVBoxLayout, - QWidget) +from PySide6.QtWidgets import (QApplication, QComboBox, QDialog, # pylint: disable=E0611 + QFileDialog, QHBoxLayout, QLabel, QLineEdit, + QListWidget, QListWidgetItem, QMainWindow, + QMessageBox, QProgressBar, QPushButton, + QSplitter, QTextEdit, QVBoxLayout, QWidget) from scapy.all import ARP, Ether, get_if_addr, srp # pylint: disable=E0611 from core import db import core.networking as net from core import vendor from core.mitm import MitmThread -from core.ollama_analyst import OllamaThread +from core.ollama_analyst import OllamaThread, fetch_ollama_models from core.platform import get_os from ui.ui_arpscan import Ui_DeviceDiscovery @@ -37,6 +37,82 @@ _RESOLVE_WORKERS = 20 +class LlmAnalysisWindow(QMainWindow): + """Standalone window that streams LLM analysis of a captured packet.""" + + def __init__(self, packet_summary: str, packet_text: str, parent=None): + super().__init__(parent) + self.setWindowTitle(f"LLM Analysis — {packet_summary}") + self.resize(780, 600) + + mono = QFont() + mono.setStyleHint(QFont.StyleHint.Monospace) + mono.setPointSize(10) + + # Packet detail (collapsed by default via a small fixed height) + self._packet_detail = QTextEdit() + self._packet_detail.setReadOnly(True) + self._packet_detail.setFont(mono) + self._packet_detail.setStyleSheet("background:#111; color:#aaa;") + self._packet_detail.setPlainText(packet_text) + self._packet_detail.setMaximumHeight(120) + + # LLM output + self._output = QTextEdit() + self._output.setReadOnly(True) + self._output.setFont(mono) + self._output.setStyleSheet("background:black; color:white;") + self._output.setPlaceholderText("Waiting for analysis…") + + splitter = QSplitter(Qt.Vertical) + splitter.addWidget(self._packet_detail) + splitter.addWidget(self._output) + splitter.setSizes([100, 480]) + + self._copy_button = QPushButton("Copy analysis") + self._copy_button.clicked.connect(self._copy_to_clipboard) + self._copy_button.setEnabled(False) + + btn_row = QHBoxLayout() + btn_row.addStretch() + btn_row.addWidget(self._copy_button) + + layout = QVBoxLayout() + layout.addWidget(QLabel("Packet:")) + layout.addWidget(splitter) + layout.addLayout(btn_row) + + container = QWidget() + container.setLayout(layout) + self.setCentralWidget(container) + + def append_token(self, token: str): + """Append a streamed token to the output, clearing the placeholder on first call.""" + if self._output.toPlainText() == "Analysing…": + self._output.clear() + cursor = self._output.textCursor() + cursor.movePosition(cursor.MoveOperation.End) + cursor.insertText(token) + self._output.setTextCursor(cursor) + self._output.ensureCursorVisible() + + def set_analysing(self): + """Reset the output area to the 'analysing' placeholder state.""" + self._output.setPlainText("Analysing…") + self._copy_button.setEnabled(False) + + def set_error(self, msg: str): + """Display an error message in the output area.""" + self._output.setPlainText(f"Error: {msg}") + + def set_done(self): + """Mark analysis as complete and enable the copy button.""" + self._copy_button.setEnabled(True) + + def _copy_to_clipboard(self): + QApplication.clipboard().setText(self._output.toPlainText()) + + class DeviceDetailsWindow(QMainWindow): # pylint: disable=too-many-instance-attributes """Window showing detailed information about a device, with MITM controls.""" @@ -104,7 +180,46 @@ def __init__( # pylint: disable=too-many-arguments,too-many-positional-argument self._save_pcap_button.setEnabled(False) layout.addWidget(self._save_pcap_button) - mono = QFont("Monospace") + self._build_packet_panel(layout) + + self._captured_packets = [] # newest first, mirrors list order + self._ollama_thread: OllamaThread | None = None + self._llm_window: LlmAnalysisWindow | None = None + + central_widget = QWidget() + central_widget.setLayout(layout) + self.setCentralWidget(central_widget) + self.resize(680, 850) + + self._load_ollama_models() + + @Slot(bool) + def _toggle_mitm(self, checked): + if checked: + self._mitm = MitmThread(self.interface, self.ip_address, self.gateway_ip) + self._mitm.statusChanged.connect(self._on_status) + self._mitm.packetCaptured.connect(self._on_packet) + self._mitm.stopped.connect(self._on_mitm_stopped) + self._mitm.start() + self._mitm_button.setText("Stop MITM") + self._status_label.setStyleSheet("color: #ff8800; font-weight: bold") + else: + self._mitm_button.setEnabled(False) + if self._mitm: + self._mitm.stop() # async — _on_mitm_stopped will reset UI + + @Slot() + def _on_mitm_stopped(self): + self._mitm = None + self._mitm_button.setText("Start MITM") + self._mitm_button.setChecked(False) + self._mitm_button.setEnabled(True) + self._status_label.setStyleSheet("color: grey") + + def _build_packet_panel(self, layout: QVBoxLayout): + """Build the packet list/detail/LLM controls and add them to *layout*.""" + mono = QFont() + mono.setStyleHint(QFont.StyleHint.Monospace) mono.setPointSize(9) self._packet_list = QListWidget() @@ -117,19 +232,10 @@ def __init__( # pylint: disable=too-many-arguments,too-many-positional-argument self._packet_detail.setStyleSheet("background:black; color:white;") self._packet_detail.setPlaceholderText("Select a packet to inspect it...") - # Analyse button + LLM output self._analyse_button = QPushButton("Analyse with LLM") self._analyse_button.setEnabled(False) self._analyse_button.clicked.connect(self._analyse_packet) - self._llm_output = QTextEdit() - self._llm_output.setReadOnly(True) - self._llm_output.setFont(mono) - self._llm_output.setStyleSheet("background:black; color:white;") - self._llm_output.setPlaceholderText("LLM analysis will appear here...") - self._llm_output.setMaximumHeight(180) - - # Stack: packet list | packet detail | [analyse btn + llm output] pkt_splitter = QSplitter(Qt.Vertical) pkt_splitter.addWidget(self._packet_list) pkt_splitter.addWidget(self._packet_detail) @@ -141,44 +247,35 @@ def __init__( # pylint: disable=too-many-arguments,too-many-positional-argument ) self._user_context.returnPressed.connect(self._analyse_packet) + self._model_combo = QComboBox() + self._model_combo.setPlaceholderText("Select Ollama model…") + self._refresh_models_button = QPushButton("↻") + self._refresh_models_button.setFixedWidth(28) + self._refresh_models_button.setToolTip("Refresh available Ollama models") + self._refresh_models_button.clicked.connect(self._load_ollama_models) + model_row = QHBoxLayout() + model_row.addWidget(QLabel("Model:")) + model_row.addWidget(self._model_combo, stretch=1) + model_row.addWidget(self._refresh_models_button) + layout.addWidget(QLabel("Captured packets:")) layout.addWidget(pkt_splitter) layout.addWidget(QLabel("Context:")) layout.addWidget(self._user_context) + layout.addLayout(model_row) layout.addWidget(self._analyse_button) - layout.addWidget(QLabel("LLM analysis:")) - layout.addWidget(self._llm_output) - - self._captured_packets = [] # newest first, mirrors list order - self._ollama_thread: OllamaThread | None = None - - central_widget = QWidget() - central_widget.setLayout(layout) - self.setCentralWidget(central_widget) - self.resize(680, 850) - @Slot(bool) - def _toggle_mitm(self, checked): - if checked: - self._mitm = MitmThread(self.interface, self.ip_address, self.gateway_ip) - self._mitm.statusChanged.connect(self._on_status) - self._mitm.packetCaptured.connect(self._on_packet) - self._mitm.stopped.connect(self._on_mitm_stopped) - self._mitm.start() - self._mitm_button.setText("Stop MITM") - self._status_label.setStyleSheet("color: #ff8800; font-weight: bold") + def _load_ollama_models(self): + """Populate the model combo box with models available on the local Ollama server.""" + models = fetch_ollama_models() + self._model_combo.clear() + if models: + self._model_combo.addItems(models) + self._model_combo.setCurrentIndex(0) + self._model_combo.setEnabled(True) else: - self._mitm_button.setEnabled(False) - if self._mitm: - self._mitm.stop() # async — _on_mitm_stopped will reset UI - - @Slot() - def _on_mitm_stopped(self): - self._mitm = None - self._mitm_button.setText("Start MITM") - self._mitm_button.setChecked(False) - self._mitm_button.setEnabled(True) - self._status_label.setStyleSheet("color: grey") + self._model_combo.setPlaceholderText("Ollama not running or no models pulled") + self._model_combo.setEnabled(False) @Slot(str) def _on_status(self, msg): @@ -204,7 +301,6 @@ def _on_packet_selected(self, row): pkt = self._captured_packets[row] self._packet_detail.setPlainText(_format_packet(pkt)) self._analyse_button.setEnabled(True) - self._llm_output.clear() @Slot() def _analyse_packet(self): @@ -216,10 +312,22 @@ def _analyse_packet(self): if self._ollama_thread and self._ollama_thread.isRunning(): self._ollama_thread.terminate() - pkt_text = _format_packet(self._captured_packets[row]) + model = self._model_combo.currentText() + if not model: + QMessageBox.warning(self, "No model", "No Ollama model selected — click ↻ to refresh.") + return + + pkt = self._captured_packets[row] + pkt_text = _format_packet(pkt) user_context = self._user_context.text().strip() + + self._llm_window = LlmAnalysisWindow(pkt.summary(), pkt_text, parent=self) + self._llm_window.set_analysing() + self._llm_window.show() + self._ollama_thread = OllamaThread( pkt_text, + model, user_context=user_context, device_vendor=self._device_vendor, hostname=self._hostname, @@ -230,26 +338,22 @@ def _analyse_packet(self): self._ollama_thread.start() self._analyse_button.setEnabled(False) - self._llm_output.setPlainText("Analysing...") @Slot(str) def _on_llm_token(self, token): - cursor = self._llm_output.textCursor() - cursor.movePosition(cursor.MoveOperation.End) - # Clear placeholder on first real token - if self._llm_output.toPlainText() == "Analysing...": - self._llm_output.clear() - cursor.insertText(token) - self._llm_output.setTextCursor(cursor) - self._llm_output.ensureCursorVisible() + if self._llm_window: + self._llm_window.append_token(token) @Slot(str) def _on_llm_error(self, msg): - self._llm_output.setPlainText(f"Error: {msg}") + if self._llm_window: + self._llm_window.set_error(msg) @Slot() def _on_llm_finished(self): self._analyse_button.setEnabled(True) + if self._llm_window: + self._llm_window.set_done() @Slot() def _save_pcap(self): @@ -586,8 +690,6 @@ def _shutdown(self): if self.arp_scanner_thread is not None and self.arp_scanner_thread.isRunning(): self.arp_scanner_thread.quit() self.arp_scanner_thread.wait() - from PySide6.QtWidgets import QApplication # pylint: disable=E0611 - QApplication.quit() @@ -615,6 +717,12 @@ def __init__(self, interface, mac_vendor_lookup, timeout=1, target_cidr=None): def run(self): """Determine the target network and start the ARP scan.""" src_ip = get_if_addr(self.interface) + print(f"[scan] interface={self.interface!r} src_ip={src_ip} use_native={self.use_native}") + + if src_ip == "0.0.0.0": + print(f"[scan] ERROR: could not get a valid IP for interface {self.interface!r}") + self.finished.emit([]) + return if self.target_cidr: try: @@ -640,7 +748,10 @@ def run(self): def _scan_network(self, src_ip, network): hosts = [str(ip) for ip in network.hosts() if str(ip) != src_ip] if self.use_native: - print("Using native ARP scanner") + print( + f"[scan] native ARP scanner — {len(hosts)} hosts" + f" timeout={self.timeout}s workers={_RESOLVE_WORKERS}" + ) return self._run_native_scan(src_ip, hosts) print(f"Using Scapy ARP scanner — {len(hosts)} hosts") return self._run_scapy_scan(hosts) @@ -650,29 +761,37 @@ def _scan_network(self, src_ip, network): # ------------------------------------------------------------------ def _run_native_scan(self, src_ip, hosts): - """Scan hosts sequentially using the native C arpscanner extension.""" - arp_results = [] + """Scan hosts in parallel using the native C arpscanner extension.""" total = len(hosts) - for count, ip in enumerate(hosts, start=1): + completed = 0 + arp_results = [] + timeout_ms = int(self.timeout * 1000) + + def scan_one(ip): try: result = arpscanner.perform_arp_scan( - self.interface, - str(src_ip), - str(ip), - int(self.timeout * 1000), + self.interface, str(src_ip), str(ip), timeout_ms ) + if result: + print(f"[scan] found {ip} -> {result.get('mac')}") + return ip, result except Exception as e: # pylint: disable=broad-exception-caught - print(f"Error scanning {ip}: {e}") - result = None - - if result: - mac = result["mac"] - device_vendor = self.mac_vendor_lookup.lookup_vendor(mac) - hostname = net.get_hostname(ip) - arp_response = _make_arp_response(ip, mac) - arp_results.append((ip, mac, hostname, device_vendor, arp_response)) + print(f"[scan] error scanning {ip}: {e}") + return ip, None + + with concurrent.futures.ThreadPoolExecutor(max_workers=_RESOLVE_WORKERS) as pool: + for ip, result in pool.map(scan_one, hosts): + completed += 1 + if result: + mac = result["mac"] + device_vendor = self.mac_vendor_lookup.lookup_vendor(mac) + hostname = net.get_hostname(ip) + arp_results.append( + (ip, mac, hostname, device_vendor, _make_arp_response(ip, mac)) + ) + self.partialResults.emit(list(arp_results)) + self._update_progress(completed, total, arp_results) - self._update_progress(count, total, arp_results) return arp_results # ------------------------------------------------------------------ diff --git a/core/mitm.py b/core/mitm.py index c7972d9..246dad4 100644 --- a/core/mitm.py +++ b/core/mitm.py @@ -282,11 +282,15 @@ def run(self): """Sniff packets from/to the target IP until stopped.""" self._running = True bpf = f"host {self.target_ip} and not arp" + + def _emit(pkt): + self.packetCaptured.emit(pkt) + while self._running: sniff( iface=self.interface, filter=bpf, - prn=self.packetCaptured.emit, + prn=_emit, stop_filter=lambda _: not self._running, store=False, timeout=1, diff --git a/core/ollama_analyst.py b/core/ollama_analyst.py index 2b4e16b..fc7f8fd 100644 --- a/core/ollama_analyst.py +++ b/core/ollama_analyst.py @@ -7,8 +7,18 @@ import requests from PySide6.QtCore import QThread, Signal # pylint: disable=E0611 -OLLAMA_URL = "http://localhost:11434/api/generate" -DEFAULT_MODEL = "llama3.2:1b" +OLLAMA_BASE = "http://localhost:11434" +OLLAMA_URL = f"{OLLAMA_BASE}/api/generate" + + +def fetch_ollama_models() -> list[str]: + """Return the list of model names available on the local Ollama server.""" + try: + resp = requests.get(f"{OLLAMA_BASE}/api/tags", timeout=3) + resp.raise_for_status() + return [m["name"] for m in resp.json().get("models", [])] + except Exception: # pylint: disable=broad-exception-caught + return [] SYSTEM_PROMPT = """You are an IoT security researcher specialising in vulnerability discovery on embedded and smart devices. @@ -34,10 +44,10 @@ class OllamaThread(QThread): def __init__( self, packet_text: str, + model: str, user_context: str = "", device_vendor: str = "", hostname: str = "", - model: str = DEFAULT_MODEL, parent=None, ): super().__init__(parent) diff --git a/requirements.txt b/requirements.txt index 07a8b1c..9310eb6 100644 --- a/requirements.txt +++ b/requirements.txt @@ -8,12 +8,12 @@ mccabe==0.7.0 netifaces==0.11.0 platformdirs==4.3.6 pylint==3.3.5 -PySide6==6.8.2.1 -PySide6_Addons==6.8.2.1 -PySide6_Essentials==6.8.2.1 +PySide6==6.10.1 +PySide6_Addons==6.10.1 +PySide6_Essentials==6.10.1 requests==2.33.0 scapy==2.6.1 setuptools==78.1.1 -shiboken6==6.8.2.1 +shiboken6==6.10.1 tomlkit==0.13.2 urllib3==2.6.3