From 9fad21662bfa21515958a56c9e6b7b42f86a0a1d Mon Sep 17 00:00:00 2001 From: Marco Date: Wed, 4 Mar 2026 12:56:00 +0100 Subject: [PATCH 1/2] Engine rework --- frontend/Assets/app.js | 41 ++-- leaklens.py | 82 ++----- requirements.txt | 1 - scanner/db.py | 208 +++++++++++++++++ scanner/engine.py | 491 ++++++++++++++++++++++------------------- tests/test_db.py | 233 +++++++++++++++++++ 6 files changed, 735 insertions(+), 321 deletions(-) create mode 100644 scanner/db.py create mode 100644 tests/test_db.py diff --git a/frontend/Assets/app.js b/frontend/Assets/app.js index 83df412..2ef6850 100644 --- a/frontend/Assets/app.js +++ b/frontend/Assets/app.js @@ -1,5 +1,19 @@ const API = '/api'; + // ── UNC / host helpers ───────────────────────────────────────────────────── + function parseUncHost(path) { + return path.replace(/^\\\\|^\/\//, '').split(/[\\\/]/)[0]; + } + + function parseHostPort(str) { + const ci = str.lastIndexOf(':'); + if (ci > 0) { + const p = parseInt(str.slice(ci + 1), 10); + if (p > 0 && p < 65536) return { host: str.slice(0, ci), port: p }; + } + return { host: str, port: 445 }; + } + let findings = []; let filteredFindings = []; // filtered + sorted subset rendered by virtual scroll let activeFilter = 'ALL'; @@ -17,7 +31,7 @@ // Pre-fill host from current scanPath if it looks like a UNC path const path = document.getElementById('scanPath').value.trim(); if ((path.startsWith('\\\\') || path.startsWith('//')) && !document.getElementById('smbHost').value) { - const host = path.replace(/^\\\\|^\/\//, '').split(/[\\\/]/)[0]; + const host = parseUncHost(path); if (host) document.getElementById('smbHost').value = host; } } @@ -32,7 +46,7 @@ const v = document.getElementById('scanPath').value; const isUnc = v.startsWith('\\\\') || v.startsWith('//'); if (isUnc) { - const host = v.replace(/^\\\\|^\/\//, '').split(/[\\\/]/)[0]; + const host = parseUncHost(v); const hostEl = document.getElementById('smbHost'); if (host && !hostEl.value) hostEl.value = host; } @@ -43,7 +57,7 @@ let raw = document.getElementById('smbHost').value.trim(); // If user pasted a full UNC path, extract just the host if (raw.startsWith('\\\\') || raw.startsWith('//')) { - raw = raw.replace(/^\\\\|^\/\//, '').split(/[\\\/]/)[0]; + raw = parseUncHost(raw); document.getElementById('smbHost').value = raw; } document.getElementById('sharesList').innerHTML = ''; @@ -86,13 +100,8 @@ const domain = document.getElementById('smbDomain')?.value.trim() || ''; // Parse port from the Server/Host field (e.g. "127.0.0.1:4445") - let smbPort = 445; const rawHost = document.getElementById('smbHost')?.value.trim() || ''; - const ci = rawHost.lastIndexOf(':'); - if (ci > 0) { - const p = parseInt(rawHost.slice(ci + 1), 10); - if (p > 0 && p < 65536) smbPort = p; - } + const { port: smbPort } = parseHostPort(rawHost); if (!scanPath) { alert('Please enter a file share path.'); return; } @@ -225,21 +234,13 @@ // ── SMB: Discover shares ─────────────────────────────────────────────────── async function discoverShares() { let raw = document.getElementById('smbHost').value.trim() - || (() => { - const p = document.getElementById('scanPath').value.trim(); - return p.replace(/^\\\\|^\/\//, '').split(/[\\\/]/)[0]; - })(); + || parseUncHost(document.getElementById('scanPath').value.trim()); // Normalise if someone pasted a full UNC path into the host field if (raw.startsWith('\\\\') || raw.startsWith('//')) { - raw = raw.replace(/^\\\\|^\/\//, '').split(/[\\\/]/)[0]; + raw = parseUncHost(raw); } // Parse optional :port suffix (e.g. 127.0.0.1:4445) - let host = raw, port = 445; - const colonIdx = raw.lastIndexOf(':'); - if (colonIdx > 0) { - const maybePort = parseInt(raw.slice(colonIdx + 1), 10); - if (maybePort > 0 && maybePort < 65536) { port = maybePort; host = raw.slice(0, colonIdx); } - } + const { host, port } = parseHostPort(raw); if (!host) { alert('Enter a server hostname or IP first.'); return; } const btn = document.getElementById('discoverBtn'); diff --git a/leaklens.py b/leaklens.py index fd7f903..97707f8 100644 --- a/leaklens.py +++ b/leaklens.py @@ -14,12 +14,13 @@ import os import queue import re -import sqlite3 import threading __version__ = "1.1.0" from flask import Flask, Response, jsonify, request, send_from_directory, stream_with_context +from scanner.engine import scan_path +from scanner import db as _db logging.basicConfig( level=logging.WARNING, @@ -72,14 +73,14 @@ def scan(): return jsonify({"error": "A scan is already running."}), 409 data = request.get_json(silent=True) or {} - scan_path = data.get("scanPath", "").strip() + scan_path_val = data.get("scanPath", "").strip() max_file_size_mb = int(data.get("maxFileSizeMB", 10)) workers = max(1, min(16, int(data.get("workers", 8)))) resume = bool(data.get("resume", False)) smb_port = max(1, min(65535, int(data.get("smbPort", 445)))) username, password, domain = _extract_smb_creds(data) - if not scan_path: + if not scan_path_val: return jsonify({"error": "scanPath is required."}), 400 max_size = max_file_size_mb * 1024 * 1024 @@ -93,9 +94,8 @@ def generate(): def run(): try: - from scanner.engine import scan_path as do_scan - for event in do_scan( - scan_path, + for event in scan_path( + scan_path_val, max_size, stop_event, username=username, @@ -194,24 +194,7 @@ def list_shares(): @app.route("/api/scans", methods=["GET"]) def list_scans(): """Return metadata for all completed/in-progress SQLite scan databases.""" - if not os.path.isdir(REPORTS_DIR): - return jsonify([]) - scans = [] - for fname in os.listdir(REPORTS_DIR): - if not fname.endswith(".db"): - continue - db_path = os.path.join(REPORTS_DIR, fname) - try: - conn = sqlite3.connect(f"file:{db_path}?mode=ro", uri=True) - conn.row_factory = sqlite3.Row - row = conn.execute("SELECT * FROM scans LIMIT 1").fetchone() - if row: - scans.append(dict(row)) - conn.close() - except Exception as e: - _log.debug("Could not read scan DB %s: %s", db_path, e) - scans.sort(key=lambda s: s.get("started_at", 0), reverse=True) - return jsonify(scans) + return jsonify(_db.get_all_scans(REPORTS_DIR)) # ─── GET /api/findings ──────────────────────────────────────────────────────── @@ -248,34 +231,7 @@ def get_findings(): return jsonify({"error": "Scan not found"}), 404 try: - conn = sqlite3.connect(f"file:{db_path}?mode=ro", uri=True) - conn.row_factory = sqlite3.Row - - conditions = ["scan_id = ?"] - params: list = [scan_id] - - if risk != "ALL": - conditions.append("risk_level = ?") - params.append(risk) - - if search: - conditions.append("(file_name LIKE ? OR full_path LIKE ?)") - params.extend([f"%{search}%", f"%{search}%"]) - - where = " AND ".join(conditions) - - total = conn.execute( - f"SELECT COUNT(*) FROM findings WHERE {where}", params - ).fetchone()[0] - - rows = conn.execute( - f"SELECT data FROM findings WHERE {where} ORDER BY id LIMIT ? OFFSET ?", - params + [per_page, page * per_page], - ).fetchall() - - findings = [json.loads(r["data"]) for r in rows] - conn.close() - + total, findings = _db.query_findings(db_path, scan_id, page, per_page, risk, search) return jsonify({ "scan_id": scan_id, "total": total, @@ -301,24 +257,16 @@ def export_scan(scan_id): return jsonify({"error": "Scan not found"}), 404 try: - conn = sqlite3.connect(f"file:{db_path}?mode=ro", uri=True) - conn.row_factory = sqlite3.Row - scan_row = conn.execute("SELECT * FROM scans WHERE id = ?", (scan_id,)).fetchone() - if not scan_row: - conn.close() + scan_meta = _db.get_scan_meta(db_path, scan_id) + if not scan_meta: return jsonify({"error": "Scan not found"}), 404 - rows = conn.execute( - "SELECT data FROM findings WHERE scan_id = ? ORDER BY id", - (scan_id,), - ).fetchall() - findings = [json.loads(r["data"]) for r in rows] - conn.close() + findings = _db.get_all_findings(db_path, scan_id) return jsonify({ "scan_id": scan_id, - "scan_path": scan_row["scan_path"], - "scan_date": scan_row["scan_date"], - "scanned": scan_row["scanned"], - "hits": scan_row["hits"], + "scan_path": scan_meta["scan_path"], + "scan_date": scan_meta["scan_date"], + "scanned": scan_meta["scanned"], + "hits": scan_meta["hits"], "findings": findings, }) except Exception: diff --git a/requirements.txt b/requirements.txt index 681be18..f8b7b0d 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,3 +1,2 @@ flask==3.1.3 smbprotocol==1.16.0 -PyYAML==6.0.2 diff --git a/scanner/db.py b/scanner/db.py new file mode 100644 index 0000000..2ac7ce8 --- /dev/null +++ b/scanner/db.py @@ -0,0 +1,208 @@ +""" +LeakLens database helpers — SQLite access layer. + +All scan data lives in per-scan databases at: + /LeakLens_.db + +Schema is applied at connection time by open_db(). Query functions open +read-only URI connections so they are safe to call concurrently with an +active scan writing to the same file. +""" +import datetime +import json +import logging +import os +import sqlite3 +import time + +_log = logging.getLogger("leaklens.db") + +# ─── Schema ─────────────────────────────────────────────────────────────────── + +DB_SCHEMA = """ +PRAGMA journal_mode=WAL; + +CREATE TABLE IF NOT EXISTS scans ( + id TEXT PRIMARY KEY, + scan_path TEXT, + scan_date TEXT, + scanned INTEGER DEFAULT 0, + hits INTEGER DEFAULT 0, + completed INTEGER DEFAULT 0, + started_at INTEGER +); + +CREATE TABLE IF NOT EXISTS findings ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + scan_id TEXT NOT NULL, + risk_level TEXT NOT NULL, + confidence INTEGER, + file_name TEXT, + full_path TEXT, + data TEXT NOT NULL +); + +CREATE INDEX IF NOT EXISTS idx_findings_scan + ON findings (scan_id, risk_level); +""" + + +# ─── Connection helpers ─────────────────────────────────────────────────────── + +def open_db(path: str) -> sqlite3.Connection: + """ + Open (or create) a scan database, apply schema, and return the connection. + Caller is responsible for calling conn.close(). + """ + conn = sqlite3.connect(path) + conn.executescript(DB_SCHEMA) + return conn + + +def open_db_readonly(path: str) -> sqlite3.Connection: + """Open an existing scan database in read-only mode.""" + conn = sqlite3.connect(f"file:{path}?mode=ro", uri=True) + conn.row_factory = sqlite3.Row + return conn + + +# ─── Write helpers (used by engine.py) ─────────────────────────────────────── + +def insert_scan(conn: sqlite3.Connection, scan_id: str, scan_path: str, + started_at: int) -> None: + """Insert the initial scan record and commit.""" + conn.execute( + "INSERT OR REPLACE INTO scans (id, scan_path, scan_date, started_at) " + "VALUES (?, ?, ?, ?)", + ( + scan_id, + scan_path, + datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S"), + started_at, + ), + ) + conn.commit() + + +def update_scan_complete(conn: sqlite3.Connection, scan_id: str, + scanned: int, hits: int) -> None: + """Mark a scan as completed and record final file/hit counts.""" + conn.execute( + "UPDATE scans SET scanned=?, hits=?, completed=1 WHERE id=?", + (scanned, hits, scan_id), + ) + conn.commit() + + +# ─── Query helpers (used by leaklens.py route handlers) ────────────────────── + +def query_findings( + db_path: str, + scan_id: str, + page: int, + per_page: int, + risk: str, + search: str, +) -> tuple: + """ + Query paginated findings from a scan database. + + Returns (total: int, findings: list[dict]). + Raises FileNotFoundError if db_path does not exist. + Raises sqlite3.Error on database errors. + """ + if not os.path.isfile(db_path): + raise FileNotFoundError(f"Scan database not found: {db_path}") + + conn = open_db_readonly(db_path) + try: + conditions = ["scan_id = ?"] + params: list = [scan_id] + + if risk != "ALL": + conditions.append("risk_level = ?") + params.append(risk) + + if search: + conditions.append("(file_name LIKE ? OR full_path LIKE ?)") + params.extend([f"%{search}%", f"%{search}%"]) + + where = " AND ".join(conditions) + + total = conn.execute( + f"SELECT COUNT(*) FROM findings WHERE {where}", params + ).fetchone()[0] + + rows = conn.execute( + f"SELECT data FROM findings WHERE {where} ORDER BY id LIMIT ? OFFSET ?", + params + [per_page, page * per_page], + ).fetchall() + + findings = [json.loads(r["data"]) for r in rows] + return total, findings + finally: + conn.close() + + +def get_scan_meta(db_path: str, scan_id: str) -> dict | None: + """ + Return the scan metadata row for scan_id as a dict, or None if not found. + Raises FileNotFoundError if db_path does not exist. + """ + if not os.path.isfile(db_path): + raise FileNotFoundError(f"Scan database not found: {db_path}") + + conn = open_db_readonly(db_path) + try: + row = conn.execute( + "SELECT * FROM scans WHERE id = ?", (scan_id,) + ).fetchone() + return dict(row) if row else None + finally: + conn.close() + + +def get_all_findings(db_path: str, scan_id: str) -> list: + """ + Return all findings for a scan as a list of dicts. + Raises FileNotFoundError if db_path does not exist. + """ + if not os.path.isfile(db_path): + raise FileNotFoundError(f"Scan database not found: {db_path}") + + conn = open_db_readonly(db_path) + try: + rows = conn.execute( + "SELECT data FROM findings WHERE scan_id = ? ORDER BY id", + (scan_id,), + ).fetchall() + return [json.loads(r["data"]) for r in rows] + finally: + conn.close() + + +def get_all_scans(reports_dir: str) -> list: + """ + Read scan metadata from all LeakLens_*.db files in reports_dir. + Returns list of scan dicts sorted by started_at descending. + Non-readable databases are skipped with a debug log. + """ + if not os.path.isdir(reports_dir): + return [] + + scans = [] + for fname in os.listdir(reports_dir): + if not fname.startswith("LeakLens_") or not fname.endswith(".db"): + continue + db_path = os.path.join(reports_dir, fname) + try: + conn = open_db_readonly(db_path) + row = conn.execute("SELECT * FROM scans LIMIT 1").fetchone() + if row: + scans.append(dict(row)) + conn.close() + except Exception as e: + _log.debug("Could not read scan DB %s: %s", db_path, e) + + scans.sort(key=lambda s: s.get("started_at", 0), reverse=True) + return scans diff --git a/scanner/engine.py b/scanner/engine.py index 7ac81c9..ca9e645 100644 --- a/scanner/engine.py +++ b/scanner/engine.py @@ -7,10 +7,10 @@ import datetime import json import queue -import sqlite3 import time import threading -from concurrent.futures import ThreadPoolExecutor, wait as fut_wait +from concurrent.futures import ThreadPoolExecutor +from dataclasses import dataclass from typing import Generator from scanner.patterns import ( @@ -32,6 +32,30 @@ read_smb_file, SMB_AVAILABLE, ) +from scanner import db as _db + + +# ─── Module-level constants ──────────────────────────────────────────────────── + +_CHECKPOINT_EVERY = 1000 # save checkpoint every N files processed +_DB_COMMIT_EVERY = 50 # commit SQLite batch every N inserts + + +# ─── Scan configuration ─────────────────────────────────────────────────────── + +@dataclass +class _ScanConfig: + """Immutable scan parameters shared across worker threads.""" + root: str + max_size: int + stop_event: threading.Event + is_smb: bool + smb_port: int + username: str | None = None + password: str | None = None + domain: str | None = None + suppressions: object = None # loaded once at scan start + unc_share_root: str = "" # empty string for local scans # ─── Local path scanner ─────────────────────────────────────────────────────── @@ -58,18 +82,18 @@ def _check_local_file(path: str, max_size: int): matched = [] - if binary_risk: - pass - elif ext in TARGET_EXTENSIONS or (risky_name and ext == ""): - if stat.st_size <= max_size: - try: - with open(path, "r", encoding="utf-8", errors="ignore") as f: - content = f.read() - matched = scan_content(content, path) - except OSError: - pass - elif not risky_name: - return None + # Binary-risk files (.kdbx, .pfx, …) are flagged by file type alone; no content scan. + if not binary_risk: + if ext in TARGET_EXTENSIONS or (risky_name and ext == ""): + if stat.st_size <= max_size: + try: + with open(path, "r", encoding="utf-8", errors="ignore") as f: + content = f.read() + matched = scan_content(content, path) + except OSError: + pass + elif not risky_name: + return None if not binary_risk and not risky_name and not matched: return None @@ -146,34 +170,193 @@ def _scan_smb( yield ("smb_file", smb_path, fname, stat, server, share) -# ─── SQLite schema ──────────────────────────────────────────────────────────── - -_DB_SCHEMA = """ -PRAGMA journal_mode=WAL; - -CREATE TABLE IF NOT EXISTS scans ( - id TEXT PRIMARY KEY, - scan_path TEXT, - scan_date TEXT, - scanned INTEGER DEFAULT 0, - hits INTEGER DEFAULT 0, - completed INTEGER DEFAULT 0, - started_at INTEGER -); - -CREATE TABLE IF NOT EXISTS findings ( - id INTEGER PRIMARY KEY AUTOINCREMENT, - scan_id TEXT NOT NULL, - risk_level TEXT NOT NULL, - confidence INTEGER, - file_name TEXT, - full_path TEXT, - data TEXT NOT NULL -); - -CREATE INDEX IF NOT EXISTS idx_findings_scan - ON findings (scan_id, risk_level); -""" +# ─── Per-file analysis ──────────────────────────────────────────────────────── + +def _analyse_item(item: tuple, cfg: _ScanConfig) -> tuple: + """ + Analyse one file item (local or SMB). + Returns (finding_dict_or_None, file_path_str). + """ + if item[0] == "file": + _, fpath = item + finding = _check_local_file(fpath, cfg.max_size) + if finding is None: + return None, fpath + pids = [f["id"] for f in finding.get("findingsDetail", [])] + if is_suppressed(fpath, pids, cfg.suppressions, cfg.root): + return None, fpath + return finding, fpath + + # "smb_file" + _, smb_path, fname, stat, _srv, _shr = item + if fname.lower() in EXCLUDED_FILENAMES: + return None, smb_path + ext = os.path.splitext(fname)[1].lower() + name_base = os.path.splitext(fname.lower())[0] + binary_risk = ext in FLAGGED_EXTENSIONS + risky_name = ( + fname.lower() in FLAGGED_EXACT_NAMES + or any(n in name_base for n in FLAGGED_NAMES) + ) + size_bytes = stat.st_size if hasattr(stat, "st_size") else 0 + matched = [] + + # Binary-risk files (.kdbx, .pfx, …) are flagged by file type alone; no content scan. + if not binary_risk: + if ext in TARGET_EXTENSIONS or (risky_name and ext == ""): + if size_bytes <= cfg.max_size: + content = read_smb_file(smb_path, cfg.max_size, port=cfg.smb_port) + if content: + matched = scan_content(content, smb_path) + elif not risky_name: + return None, smb_path + + if not binary_risk and not risky_name and not matched: + return None, smb_path + + try: + last_modified = datetime.datetime.fromtimestamp(stat.st_mtime).strftime("%Y-%m-%d %H:%M") + last_accessed = datetime.datetime.fromtimestamp(stat.st_atime).strftime("%Y-%m-%d %H:%M") + except Exception: + last_modified = last_accessed = "" + + try: + rel_path = smb_path[len(cfg.unc_share_root):].lstrip("\\") + except Exception: + rel_path = smb_path + + smb_meta = { + "smbShare": cfg.unc_share_root, + "smbServer": _srv, + "smbRelativePath": rel_path, + } + + finding = build_finding( + path=smb_path, ext=ext, size_bytes=size_bytes, + last_modified=last_modified, last_accessed=last_accessed, + owner="", matched_patterns=matched, + risky_name=risky_name, binary_risk=binary_risk, smb_meta=smb_meta, + ) + if finding is None: + return None, smb_path + + pids = [f["id"] for f in finding.get("findingsDetail", [])] + if is_suppressed(smb_path, pids, cfg.suppressions, cfg.unc_share_root): + return None, smb_path + + return finding, smb_path + + +# ─── Walk thread ────────────────────────────────────────────────────────────── + +def _walk( + cfg: _ScanConfig, + resume_from: str | None, + file_q: queue.Queue, + event_q: queue.Queue, + workers: int, +) -> None: + """ + Walk the filesystem and enqueue file items into file_q. + If resuming, skips files until resume_from path is passed. + Errors/logs go directly to event_q (never to file_q). + """ + _skipping = resume_from is not None + skipped = 0 + + try: + if cfg.is_smb: + for ev in _scan_smb( + cfg.root, cfg.max_size, cfg.stop_event, + cfg.username, cfg.password, cfg.domain, cfg.smb_port, + ): + if cfg.stop_event.is_set(): + break + if ev[0] == "error": + event_q.put({"type": "error", "message": ev[1]}) + break + if ev[0] == "log": + event_q.put({"type": "log", "message": ev[1]}) + continue + # "smb_file" — resume skip check + if _skipping: + if ev[1] == resume_from: + _skipping = False + skipped += 1 + continue + file_q.put(ev) + else: + for ev in _scan_local(cfg.root, cfg.max_size, cfg.stop_event): + if cfg.stop_event.is_set(): + break + if _skipping: + if ev[1] == resume_from: + _skipping = False + skipped += 1 + continue + file_q.put(ev) + finally: + if skipped > 0: + event_q.put({ + "type": "log", + "message": f"Resumed — skipped {skipped} already-scanned files.", + }) + for _ in range(workers): + file_q.put(None) # poison pill per worker + + +# ─── Worker ─────────────────────────────────────────────────────────────────── + +def _worker( + cfg: _ScanConfig, + file_q: queue.Queue, + event_q: queue.Queue, + lock: threading.Lock, + state: dict, +) -> None: + """Consume file_q, analyse files, emit events to event_q.""" + while not cfg.stop_event.is_set(): + try: + item = file_q.get(timeout=1.0) + except queue.Empty: + continue + + if item is None: + break # poison pill + + fpath = item[1] if len(item) > 1 else "?" + finding = None + + try: + finding, fpath = _analyse_item(item, cfg) + except Exception as e: + event_q.put({"type": "log", "message": f"[ERROR] {fpath}: {e}"}) + + # Update shared counters exactly once per file + with lock: + state["file_count"] += 1 + fc = state["file_count"] + if finding is not None: + state["hit_count"] += 1 + hc = state["hit_count"] + else: + hc = state["hit_count"] + + if finding is not None: + event_q.put(finding) + + if fc % 100 == 0 or fc == 1: + elapsed = max(0.001, time.time() - state["scan_start_time"]) + rate = round(fc / elapsed, 1) + event_q.put({ + "type": "progress", + "scanned": fc, + "hits": hc, + "current": fpath, + "rate": rate, + }) + + event_q.put(("__done__",)) # ─── Main entry point ───────────────────────────────────────────────────────── @@ -209,9 +392,6 @@ def scan_path( {type: "summary", scanned: int, hits: int, scanId: str} {type: "error", message: str} """ - _CHECKPOINT_EVERY = 1000 # checkpoint: save every N files processed - _DB_COMMIT_EVERY = 50 # SQLite: commit after N inserts - is_smb = is_unc_path(root) yield {"type": "log", "message": f"Starting scan of: {root} (workers={workers})"} @@ -222,7 +402,12 @@ def scan_path( # ── Suppression setup ───────────────────────────────────────────────────── if is_smb: - suppressions = load_suppressions(os.getcwd()) + # .leaklensignore cannot be read from the remote share. Fall back to the + # process working directory (typically the LeakLens installation root). + # Place .leaklensignore there if suppression rules are needed for SMB scans. + cwd = os.getcwd() + suppressions = load_suppressions(cwd) + yield {"type": "log", "message": f"SMB scan: .leaklensignore loaded from {cwd}"} else: if not os.path.exists(root): yield {"type": "error", "message": f"Path does not exist or is not accessible: {root}"} @@ -234,20 +419,14 @@ def scan_path( reports_dir = os.path.join(os.path.dirname(os.path.dirname(__file__)), "reports") os.makedirs(reports_dir, exist_ok=True) - timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S") - scan_id = timestamp + timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S") + scan_id = timestamp db_file = os.path.join(reports_dir, f"LeakLens_{timestamp}.db") - ckpt_file = _ckpt_path(root, reports_dir) + ckpt_file = _ckpt_path(root, reports_dir) # ── SQLite setup ────────────────────────────────────────────────────────── - _db_conn = sqlite3.connect(db_file) - _db_conn.executescript(_DB_SCHEMA) - _db_conn.execute( - "INSERT OR REPLACE INTO scans (id, scan_path, scan_date, started_at) " - "VALUES (?, ?, ?, ?)", - (scan_id, root, datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S"), int(time.time())), - ) - _db_conn.commit() + _db_conn = _db.open_db(db_file) + _db.insert_scan(_db_conn, scan_id, root, int(time.time())) _db_pending = 0 # ── Checkpoint / resume ─────────────────────────────────────────────────── @@ -260,13 +439,27 @@ def scan_path( "message": f"Resuming scan — checkpoint found, skipping up to {_ckpt_scanned} files", } - # SMB share root used for relative paths and suppression checks + # ── SMB share root used for relative paths ──────────────────────────────── if is_smb: _srv, _shr, _ = parse_unc(root) unc_share_root = f"\\\\{_srv}\\{_shr}" else: unc_share_root = "" + # ── Scan configuration (shared read-only across threads) ────────────────── + cfg = _ScanConfig( + root=root, + max_size=max_size, + stop_event=stop_event, + is_smb=is_smb, + smb_port=smb_port, + username=username, + password=password, + domain=domain, + suppressions=suppressions, + unc_share_root=unc_share_root, + ) + # ── Shared mutable state — ALL mutations guarded by _lock ───────────────── _lock = threading.Lock() _state = { @@ -279,176 +472,12 @@ def scan_path( _file_q: queue.Queue = queue.Queue(maxsize=max(workers * 8, 64)) _event_q: queue.Queue = queue.Queue(maxsize=max(workers * 32, 512)) - # ── Per-file analysis ───────────────────────────────────────────────────── - def _analyse_item(item: tuple) -> tuple: - """ - Analyse one file item (local or SMB). - Returns (finding_dict_or_None, file_path_str). - """ - if item[0] == "file": - _, fpath = item - finding = _check_local_file(fpath, max_size) - if finding is None: - return None, fpath - pids = [f["id"] for f in finding.get("findingsDetail", [])] - if is_suppressed(fpath, pids, suppressions, root): - return None, fpath - return finding, fpath - - # "smb_file" - _, smb_path, fname, stat, _srv, _shr = item - if fname.lower() in EXCLUDED_FILENAMES: - return None, smb_path - ext = os.path.splitext(fname)[1].lower() - name_base = os.path.splitext(fname.lower())[0] - binary_risk = ext in FLAGGED_EXTENSIONS - risky_name = ( - fname.lower() in FLAGGED_EXACT_NAMES - or any(n in name_base for n in FLAGGED_NAMES) - ) - size_bytes = stat.st_size if hasattr(stat, "st_size") else 0 - matched = [] - - if binary_risk: - pass - elif ext in TARGET_EXTENSIONS or (risky_name and ext == ""): - if size_bytes <= max_size: - content = read_smb_file(smb_path, max_size, port=smb_port) - if content: - matched = scan_content(content, smb_path) - elif not risky_name: - return None, smb_path - - if not binary_risk and not risky_name and not matched: - return None, smb_path - - try: - last_modified = datetime.datetime.fromtimestamp(stat.st_mtime).strftime("%Y-%m-%d %H:%M") - last_accessed = datetime.datetime.fromtimestamp(stat.st_atime).strftime("%Y-%m-%d %H:%M") - except Exception: - last_modified = last_accessed = "" - - try: - rel_path = smb_path[len(unc_share_root):].lstrip("\\") - except Exception: - rel_path = smb_path - - smb_meta = { - "smbShare": unc_share_root, - "smbServer": _srv, - "smbRelativePath": rel_path, - } - - finding = build_finding( - path=smb_path, ext=ext, size_bytes=size_bytes, - last_modified=last_modified, last_accessed=last_accessed, - owner="", matched_patterns=matched, - risky_name=risky_name, binary_risk=binary_risk, smb_meta=smb_meta, - ) - if finding is None: - return None, smb_path - - pids = [f["id"] for f in finding.get("findingsDetail", [])] - if is_suppressed(smb_path, pids, suppressions, unc_share_root): - return None, smb_path - - return finding, smb_path - - # ── Walk thread ─────────────────────────────────────────────────────────── - def _walk() -> None: - """ - Walk the filesystem and enqueue file items into _file_q. - If resuming, skips files until _resume_from path is passed. - Errors/logs go directly to _event_q (never to _file_q). - """ - _skipping = _resume_from is not None - skipped = 0 - - try: - if is_smb: - for ev in _scan_smb(root, max_size, stop_event, username, password, domain, smb_port): - if stop_event.is_set(): - break - if ev[0] == "error": - _event_q.put({"type": "error", "message": ev[1]}) - break - if ev[0] == "log": - _event_q.put({"type": "log", "message": ev[1]}) - continue - # "smb_file" — resume skip check - if _skipping: - if ev[1] == _resume_from: - _skipping = False - skipped += 1 - continue - _file_q.put(ev) - else: - for ev in _scan_local(root, max_size, stop_event): - if stop_event.is_set(): - break - if _skipping: - if ev[1] == _resume_from: - _skipping = False - skipped += 1 - continue - _file_q.put(ev) - finally: - if skipped > 0: - _event_q.put({ - "type": "log", - "message": f"Resumed — skipped {skipped} already-scanned files.", - }) - for _ in range(workers): - _file_q.put(None) # poison pill per worker - - # ── Worker ──────────────────────────────────────────────────────────────── - def _worker() -> None: - """Consume _file_q, analyse files, emit events to _event_q.""" - while not stop_event.is_set(): - try: - item = _file_q.get(timeout=1.0) - except queue.Empty: - continue - - if item is None: - break # poison pill - - fpath = item[1] if len(item) > 1 else "?" - finding = None - - try: - finding, fpath = _analyse_item(item) - except Exception as e: - _event_q.put({"type": "log", "message": f"[ERROR] {fpath}: {e}"}) - - # Update shared counters exactly once per file - with _lock: - _state["file_count"] += 1 - fc = _state["file_count"] - if finding is not None: - _state["hit_count"] += 1 - hc = _state["hit_count"] - else: - hc = _state["hit_count"] - - if finding is not None: - _event_q.put(finding) - - if fc % 100 == 0 or fc == 1: - elapsed = max(0.001, time.time() - _state["scan_start_time"]) - rate = round(fc / elapsed, 1) - _event_q.put({ - "type": "progress", - "scanned": fc, - "hits": hc, - "current": fpath, - "rate": rate, - }) - - _event_q.put(("__done__",)) - # ── Orchestrate ─────────────────────────────────────────────────────────── - walk_thread = threading.Thread(target=_walk, daemon=True) + walk_thread = threading.Thread( + target=_walk, + args=(cfg, _resume_from, _file_q, _event_q, workers), + daemon=True, + ) walk_thread.start() _ckpt_last_scanned = 0 @@ -456,7 +485,7 @@ def _worker() -> None: with ThreadPoolExecutor(max_workers=workers) as executor: for _ in range(workers): - executor.submit(_worker) + executor.submit(_worker, cfg, _file_q, _event_q, _lock, _state) done_workers = 0 fatal_error = False @@ -536,11 +565,7 @@ def _worker() -> None: # ── Finalise SQLite scan record ─────────────────────────────────────────── try: - _db_conn.execute( - "UPDATE scans SET scanned=?, hits=?, completed=1 WHERE id=?", - (fc, hc, scan_id), - ) - _db_conn.commit() + _db.update_scan_complete(_db_conn, scan_id, fc, hc) except Exception: pass diff --git a/tests/test_db.py b/tests/test_db.py new file mode 100644 index 0000000..a892da3 --- /dev/null +++ b/tests/test_db.py @@ -0,0 +1,233 @@ +""" +Tests for scanner/db.py — SQLite access layer. +""" +import json +import os +import sqlite3 +import pytest + +from scanner.db import ( + open_db, + open_db_readonly, + insert_scan, + update_scan_complete, + query_findings, + get_scan_meta, + get_all_findings, + get_all_scans, + DB_SCHEMA, +) + + +# ─── Fixtures ───────────────────────────────────────────────────────────────── + +@pytest.fixture +def db_path(tmp_path): + return str(tmp_path / "LeakLens_test.db") + + +@pytest.fixture +def conn(db_path): + c = open_db(db_path) + yield c + c.close() + + +def _seed_scan(conn, scan_id="20240101_120000", scan_path="/tmp/test"): + insert_scan(conn, scan_id, scan_path, started_at=1000) + + +def _seed_finding(conn, scan_id="20240101_120000", risk="HIGH", confidence=9, + file_name="id_rsa", full_path="/tmp/test/id_rsa"): + conn.execute( + "INSERT INTO findings (scan_id, risk_level, confidence, file_name, full_path, data) " + "VALUES (?, ?, ?, ?, ?, ?)", + (scan_id, risk, confidence, file_name, full_path, + json.dumps({"riskLevel": risk, "fileName": file_name, "fullPath": full_path})), + ) + conn.commit() + + +# ─── open_db / schema ───────────────────────────────────────────────────────── + +def test_open_db_creates_tables(db_path): + conn = open_db(db_path) + tables = {r[0] for r in conn.execute("SELECT name FROM sqlite_master WHERE type='table'").fetchall()} + conn.close() + assert "scans" in tables + assert "findings" in tables + + +def test_open_db_creates_index(db_path): + conn = open_db(db_path) + indexes = {r[0] for r in conn.execute("SELECT name FROM sqlite_master WHERE type='index'").fetchall()} + conn.close() + assert "idx_findings_scan" in indexes + + +def test_open_db_idempotent(db_path): + # Calling open_db twice should not fail (IF NOT EXISTS guards) + conn1 = open_db(db_path) + conn1.close() + conn2 = open_db(db_path) + conn2.close() + + +# ─── insert_scan / update_scan_complete ────────────────────────────────────── + +def test_insert_scan(conn, db_path): + insert_scan(conn, "scan1", "/some/path", started_at=999) + row = conn.execute("SELECT * FROM scans WHERE id = 'scan1'").fetchone() + assert row is not None + assert row[1] == "/some/path" # scan_path + assert row[6] == 999 # started_at + + +def test_update_scan_complete(conn, db_path): + insert_scan(conn, "scan2", "/path", started_at=0) + update_scan_complete(conn, "scan2", scanned=100, hits=5) + row = conn.execute("SELECT scanned, hits, completed FROM scans WHERE id = 'scan2'").fetchone() + assert row[0] == 100 + assert row[1] == 5 + assert row[2] == 1 + + +# ─── query_findings ─────────────────────────────────────────────────────────── + +def test_query_findings_basic(conn, db_path): + _seed_scan(conn) + _seed_finding(conn, risk="HIGH") + _seed_finding(conn, risk="LOW", file_name="config.json", full_path="/tmp/test/config.json") + + total, results = query_findings(db_path, "20240101_120000", page=0, per_page=100, + risk="ALL", search="") + assert total == 2 + assert len(results) == 2 + + +def test_query_findings_risk_filter(conn, db_path): + _seed_scan(conn) + _seed_finding(conn, risk="HIGH") + _seed_finding(conn, risk="LOW", file_name="other.txt", full_path="/tmp/test/other.txt") + + total, results = query_findings(db_path, "20240101_120000", page=0, per_page=100, + risk="HIGH", search="") + assert total == 1 + assert results[0]["riskLevel"] == "HIGH" + + +def test_query_findings_search_filter(conn, db_path): + _seed_scan(conn) + _seed_finding(conn, file_name="id_rsa", full_path="/tmp/test/id_rsa") + _seed_finding(conn, risk="LOW", file_name="config.json", full_path="/tmp/test/config.json") + + total, results = query_findings(db_path, "20240101_120000", page=0, per_page=100, + risk="ALL", search="id_rsa") + assert total == 1 + assert results[0]["fileName"] == "id_rsa" + + +def test_query_findings_pagination(conn, db_path): + _seed_scan(conn) + for i in range(5): + _seed_finding(conn, risk="HIGH", file_name=f"file{i}.txt", + full_path=f"/tmp/test/file{i}.txt") + + total, page0 = query_findings(db_path, "20240101_120000", page=0, per_page=2, + risk="ALL", search="") + assert total == 5 + assert len(page0) == 2 + + _, page1 = query_findings(db_path, "20240101_120000", page=1, per_page=2, + risk="ALL", search="") + assert len(page1) == 2 + + _, page2 = query_findings(db_path, "20240101_120000", page=2, per_page=2, + risk="ALL", search="") + assert len(page2) == 1 + + +def test_query_findings_file_not_found(tmp_path): + with pytest.raises(FileNotFoundError): + query_findings(str(tmp_path / "missing.db"), "scan1", 0, 100, "ALL", "") + + +# ─── get_scan_meta ──────────────────────────────────────────────────────────── + +def test_get_scan_meta_found(conn, db_path): + _seed_scan(conn) + meta = get_scan_meta(db_path, "20240101_120000") + assert meta is not None + assert meta["id"] == "20240101_120000" + assert meta["scan_path"] == "/tmp/test" + + +def test_get_scan_meta_not_found(conn, db_path): + _seed_scan(conn) + result = get_scan_meta(db_path, "nonexistent") + assert result is None + + +def test_get_scan_meta_file_not_found(tmp_path): + with pytest.raises(FileNotFoundError): + get_scan_meta(str(tmp_path / "missing.db"), "scan1") + + +# ─── get_all_findings ───────────────────────────────────────────────────────── + +def test_get_all_findings(conn, db_path): + _seed_scan(conn) + _seed_finding(conn, risk="HIGH") + _seed_finding(conn, risk="LOW", file_name="other.txt", full_path="/tmp/other.txt") + + results = get_all_findings(db_path, "20240101_120000") + assert len(results) == 2 + + +def test_get_all_findings_empty(conn, db_path): + _seed_scan(conn) + results = get_all_findings(db_path, "20240101_120000") + assert results == [] + + +# ─── get_all_scans ──────────────────────────────────────────────────────────── + +def test_get_all_scans_empty_dir(tmp_path): + result = get_all_scans(str(tmp_path)) + assert result == [] + + +def test_get_all_scans_nonexistent_dir(tmp_path): + result = get_all_scans(str(tmp_path / "does_not_exist")) + assert result == [] + + +def test_get_all_scans_multiple_dbs(tmp_path): + for scan_id, ts in [("scan_a", 100), ("scan_b", 200)]: + db_p = str(tmp_path / f"LeakLens_{scan_id}.db") + c = open_db(db_p) + insert_scan(c, scan_id, f"/path/{scan_id}", started_at=ts) + c.close() + + results = get_all_scans(str(tmp_path)) + assert len(results) == 2 + # Should be sorted descending by started_at + assert results[0]["started_at"] == 200 + assert results[1]["started_at"] == 100 + + +def test_get_all_scans_ignores_non_leaklens_dbs(tmp_path): + # Create a non-LeakLens DB — should be ignored + other = str(tmp_path / "other.db") + conn = sqlite3.connect(other) + conn.close() + + # Create a valid LeakLens DB + db_p = str(tmp_path / "LeakLens_scan1.db") + c = open_db(db_p) + insert_scan(c, "scan1", "/path", started_at=0) + c.close() + + results = get_all_scans(str(tmp_path)) + assert len(results) == 1 + assert results[0]["id"] == "scan1" From 823c6924b9d7f61a57fb1b055a8f735af2048306 Mon Sep 17 00:00:00 2001 From: Marco Date: Wed, 4 Mar 2026 13:04:06 +0100 Subject: [PATCH 2/2] SMB concurrency fix --- scanner/smb.py | 121 ++++++++++++++++++++++++++++++++++++++----------- 1 file changed, 95 insertions(+), 26 deletions(-) diff --git a/scanner/smb.py b/scanner/smb.py index f063502..53c4173 100644 --- a/scanner/smb.py +++ b/scanner/smb.py @@ -4,6 +4,8 @@ """ import re +import threading +import time try: import smbclient @@ -13,6 +15,49 @@ SMB_AVAILABLE = False +# ─── SMB concurrency limiter ────────────────────────────────────────────────── +# +# SMB2 uses a credit-based flow-control system. Each request (open, read, stat, +# close) costs 1 credit; the server replenishes credits in responses. When many +# worker threads fire simultaneous requests against the same connection the pool +# can hit 0, which causes subsequent requests — including GC-triggered Close +# requests from finalised directory handles — to raise: +# SMBException: Request requires 1 credits but only 0 credits are available +# +# _SMB_SEM caps the total number of in-flight SMB I/O operations across the +# walk thread and all worker threads. A value of 4 is conservative enough to +# prevent exhaustion against both Samba and Windows servers without significantly +# reducing throughput (directory listings are cheap; file reads are the bottleneck). + +_SMB_SEM = threading.Semaphore(4) + +# Delays (seconds) between successive retries on credit exhaustion. +_CREDIT_DELAYS = (0.05, 0.15, 0.4, 1.0, 2.5) + + +def _is_credit_error(exc: Exception) -> bool: + return "credits" in str(exc).lower() + + +def _smb_retry(fn, *args, **kwargs): + """ + Call fn(*args, **kwargs), retrying with backoff on SMB credit-exhaustion. + All other exceptions are re-raised immediately. + After all retries are exhausted the final attempt's exception propagates. + """ + for delay in _CREDIT_DELAYS: + try: + return fn(*args, **kwargs) + except Exception as e: + if _is_credit_error(e): + time.sleep(delay) + else: + raise + return fn(*args, **kwargs) # final attempt — let any exception propagate + + +# ─── Path utilities ─────────────────────────────────────────────────────────── + def is_unc_path(path: str) -> bool: r"""Return True if path is a UNC path (\\server\share or //server/share).""" return path.startswith("\\\\") or path.startswith("//") @@ -36,6 +81,8 @@ def parse_unc(path: str) -> tuple[str, str, str]: return server, share, rel +# ─── Session management ─────────────────────────────────────────────────────── + def register_session(server: str, username: str = None, password: str = None, domain: str = None, port: int = 445): """Register an SMB session for the given server.""" @@ -70,6 +117,8 @@ def register_session(server: str, username: str = None, password: str = None, smbclient.register_session(server, **kwargs) +# ─── Share enumeration ──────────────────────────────────────────────────────── + def list_shares(host: str, username: str = None, password: str = None, domain: str = None, port: int = 445) -> list[str]: """ @@ -132,6 +181,8 @@ def list_shares(host: str, username: str = None, password: str = None, return shares +# ─── Directory walk ─────────────────────────────────────────────────────────── + def walk_smb(unc_root: str, stop_event=None, port: int = 445): """ Recursively walk an SMB share. @@ -145,43 +196,61 @@ def walk_smb(unc_root: str, stop_event=None, port: int = 445): def _recurse(path): if stop_event and stop_event.is_set(): return + + # Consume the entire scandir listing immediately under the semaphore so + # the SMBDirectoryIO handle is closed before we yield anything to the + # caller. Leaving the handle open inside a paused generator causes GC + # finalizer errors ("0 credits") when concurrent worker threads have + # exhausted the SMB credit pool. try: - for entry in smbclient.scandir(path, port=port): - if stop_event and stop_event.is_set(): - return - full = f"{path}\\{entry.name}" - if entry.is_dir(follow_symlinks=False): - yield from _recurse(full) - elif entry.is_file(follow_symlinks=False): - try: - # entry.stat() lazily fetches on default port 445; - # use smbclient.stat() with the explicit port instead. - stat = smbclient.stat(full, port=port) - yield full, entry.name, stat - except Exception as e: - yield ("__error__", full, str(e)) + with _SMB_SEM: + entries = _smb_retry( + lambda: list(smbclient.scandir(path, port=port)) + ) except Exception as e: yield ("__error__", path, str(e)) + return + + for entry in entries: + if stop_event and stop_event.is_set(): + return + full = f"{path}\\{entry.name}" + if entry.is_dir(follow_symlinks=False): + yield from _recurse(full) + elif entry.is_file(follow_symlinks=False): + try: + with _SMB_SEM: + stat = _smb_retry(smbclient.stat, full, port=port) + yield full, entry.name, stat + except Exception as e: + yield ("__error__", full, str(e)) yield from _recurse(normalize_unc(unc_root)) +# ─── File reading ───────────────────────────────────────────────────────────── + def read_smb_file(smb_path: str, max_size: int, port: int = 445) -> str | None: """Read the content of a file over SMB. Returns None on error or if too large.""" if not SMB_AVAILABLE: return None - try: - with smbclient.open_file(smb_path, mode="rb", port=port) as f: - # Check size first - f.seek(0, 2) - size = f.tell() - if size > max_size: - return None - f.seek(0) - raw = f.read(max_size) - return raw.decode("utf-8", errors="ignore") - except Exception: - return None + for delay in (*_CREDIT_DELAYS, None): + try: + with _SMB_SEM: + with smbclient.open_file(smb_path, mode="rb", port=port) as f: + f.seek(0, 2) + size = f.tell() + if size > max_size: + return None + f.seek(0) + raw = f.read(max_size) + return raw.decode("utf-8", errors="ignore") + except Exception as e: + if delay is not None and _is_credit_error(e): + time.sleep(delay) + continue + return None + return None def smb_file_size(smb_path: str) -> int: