driverclient 0.2.0__py3-none-any.whl

driverclient/ops/capture.py

@@ -0,0 +1,679 @@
+"""
+client/ops/capture.py — Three capture operations for the Driver Server client.
+
+capture_all()       Dump every package in DriverStore to the repo.
+                    Server deduplicates by SHA. No resolve needed.
+
+capture_missing()   Scan + resolve to find HWIDs the repo lacks, then export
+                    any of those that are already installed locally and submit
+                    them.  No Windows Update.
+
+wu_update()         Scan + resolve for not-found HWIDs, trigger Windows Update,
+                    poll incrementally (submit each new batch as it appears),
+                    stop after wu_wait_minutes or 2 empty polls.
+
+Upload pipeline (all three operations):
+    export pool  (parallel_exports workers) — pnputil /export-driver + pack
+         ↓  queue
+    upload pool  (parallel_uploads workers) — POST /submit  (8 workers default)
+
+Public API:
+    capture_all()               -> CaptureResult
+    capture_missing(force=False) -> CaptureResult
+    wu_update(force=False)       -> CaptureResult
+"""
+import queue
+import shutil
+import subprocess
+import tempfile
+import threading
+import time
+from concurrent.futures import ThreadPoolExecutor
+from dataclasses import dataclass, field
+from datetime import datetime, timedelta, timezone
+from pathlib import Path
+
+import functools
+
+from driverclient import events
+from driverclient.config import get
+from driverclient.core.hardware import compute_fingerprint, get_model_profile
+from driverclient.core.http import BadRequestError, http_post_retry, http_upload_package, pack_driver_archive
+from driverclient.ops.resolve import ResolveResult, resolve
+from driverclient.ops.scan import Device, ScanResult, scan
+
+
+def _post_wu_session(cfg: dict, wu_start, hwids_tried: list, captured: int) -> None:
+    """Report a Windows Update session so server WU analytics has data.
+
+    Best-effort — never blocks or raises into the capture flow.
+    """
+    try:
+        duration = max(0, int((datetime.now(timezone.utc) - wu_start).total_seconds()))
+        http_post_retry(
+            f"{cfg['local_repo_url']}/submit/wu-session",
+            {
+                "machine_id": _machine_meta().get("machine_id", ""),
+                "ran_at": wu_start.isoformat(),
+                "hwids_tried": hwids_tried,
+                "wu_installed_count": captured,
+                "captured_count": captured,
+                "duration_seconds": duration,
+            },
+            timeout=cfg.get("timeout_short", 30),
+        )
+    except Exception as e:
+        events.warn("windows_update", f"  [warn] wu-session report failed: {e}",
+                    error=str(e))
+
+
+@functools.lru_cache(maxsize=1)
+def _machine_meta() -> dict:
+    """Detect this machine's vendor/model/serial/chassis + fingerprint once.
+
+    Threaded through every capture upload so the server can populate
+    submitter_* fields — otherwise the inventory shows model as 'unknown'.
+    """
+    try:
+        prof = get_model_profile()
+    except Exception:
+        prof = {}
+    try:
+        machine_id = compute_fingerprint()
+    except Exception:
+        machine_id = ""
+    return {
+        "vendor":         prof.get("vendor", ""),
+        "model":          prof.get("model", ""),
+        "friendly_model": prof.get("friendly_model", ""),
+        "serial":         prof.get("serial", ""),
+        "chassis":        prof.get("chassis_type", ""),
+        "machine_id":     machine_id,
+    }
+
+_WINDOWS_INF = Path("C:\\Windows\\INF")
+
+_PNPUTIL_KEYS: dict[str, str] = {
+    "published name": "published_name",
+    "original name":  "original_name",
+    "provider name":  "provider_name",
+    "class name":     "class_name",
+    "driver version": "driver_version",
+}
+
+
+@dataclass
+class CaptureResult:
+    submitted:     int  = 0
+    skipped:       int  = 0
+    failed:        int  = 0
+    still_missing: int  = 0
+    wu_triggered:  bool = False
+
+    def merge(self, other: "CaptureResult") -> None:
+        self.submitted += other.submitted
+        self.skipped   += other.skipped
+        self.failed    += other.failed
+
+
+# ══ Public operations ══════════════════════════════════════════════════════════
+
+def capture_all() -> CaptureResult:
+    """Export every DriverStore package and submit. Server deduplicates."""
+    cfg = get()
+    events.start("capture", "[capture-all] Enumerating DriverStore packages…")
+    packages = _enumerate_all_packages(cfg)
+    if not packages:
+        events.done("capture", "[capture-all] No packages found", total=0)
+        return CaptureResult()
+    events.progress("capture",
+                    f"[capture-all] {len(packages)} package(s) — exporting + submitting…",
+                    total=len(packages))
+    result = _export_and_submit(packages, cfg)
+    events.done("capture",
+                f"[capture-all] Done — {result.submitted} submitted  "
+                f"{result.skipped} skipped  {result.failed} failed",
+                total=len(packages), submitted=result.submitted,
+                skipped=result.skipped, failed=result.failed)
+    return result
+
+
+def capture_missing(force: bool = False) -> CaptureResult:
+    """
+    Find HWIDs the repo has no package for, export + submit any that are
+    already installed locally.  No Windows Update.
+    """
+    cfg         = get()
+    scan_result = scan(force=force)
+    res_result  = resolve(force=force)
+
+    if not res_result.has_missing:
+        events.done("capture",
+                    "[capture-missing] Repo is complete for this machine — nothing to submit",
+                    total=0)
+        return CaptureResult()
+
+    not_found_set = set(res_result.not_found)
+    events.progress("capture",
+                    f"[capture-missing] {len(not_found_set)} HWID(s) missing from repo",
+                    total=len(not_found_set))
+
+    to_export = _find_installed_packages_for_hwids(not_found_set, scan_result, cfg)
+    if not to_export:
+        events.done("capture",
+                    "[capture-missing] None of the missing HWIDs have locally-installed drivers",
+                    total=0, still_missing=len(not_found_set))
+        return CaptureResult(still_missing=len(not_found_set))
+
+    events.progress("capture",
+                    f"[capture-missing] Exporting {len(to_export)} locally-installed package(s)…",
+                    total=len(to_export))
+    result = _export_and_submit(to_export, cfg)
+    result.still_missing = max(0, len(not_found_set) - result.submitted)
+    events.done("capture",
+                f"[capture-missing] Done — {result.submitted} submitted  "
+                f"{result.skipped} skipped  {result.failed} failed",
+                total=len(to_export), submitted=result.submitted,
+                skipped=result.skipped, failed=result.failed,
+                still_missing=result.still_missing)
+    return result
+
+
+def wu_update(force: bool = False) -> CaptureResult:
+    """
+    Scan → resolve → filter not-found HWIDs to those with no local driver
+    and not in wu_ignore_classes → trigger WU → poll → submit new packages.
+
+    Gate logic (applied before WU is triggered):
+      has_driver=True   → skip (driver on machine; capture_missing handles repo upload)
+      class in ignore   → skip (wu_ignore_classes config)
+      not in scan       → include (unknown device — safer to attempt)
+    """
+    cfg         = get()
+    scan_result = scan(force=force)
+    res_result  = resolve(force=force)
+
+    if not res_result.has_missing:
+        events.done("windows_update", "[wu-update] No missing HWIDs — nothing to trigger WU for",
+                    total=0)
+        return CaptureResult()
+
+    wu_hwids = _filter_wu_hwids(res_result.not_found, scan_result, cfg)
+    skipped  = len(res_result.not_found) - len(wu_hwids)
+
+    if not wu_hwids:
+        events.done("windows_update",
+                    "[wu-update] All missing HWIDs have local drivers or are ignored — skipping WU",
+                    total=0, still_missing=len(res_result.not_found))
+        return CaptureResult(still_missing=len(res_result.not_found))
+
+    if skipped:
+        events.start("windows_update",
+                     f"[wu-update] {len(wu_hwids)} HWID(s) need WU  "
+                     f"({skipped} skipped — has local driver or ignored class)",
+                     total=len(wu_hwids), skipped=skipped)
+    else:
+        events.start("windows_update",
+                     f"[wu-update] {len(wu_hwids)} HWID(s) not in repo — triggering Windows Update…",
+                     total=len(wu_hwids))
+
+    if not _trigger_wu():
+        return CaptureResult(still_missing=len(wu_hwids))
+
+    wu_start = datetime.now(timezone.utc)
+    total    = _poll_and_submit_incremental(wu_start, cfg)
+    total.wu_triggered  = True
+    total.still_missing = max(0, len(wu_hwids) - total.submitted)
+    _post_wu_session(cfg, wu_start, list(wu_hwids), total.submitted)
+    events.done("windows_update",
+                f"[wu-update] Done — {total.submitted} captured  "
+                f"{total.still_missing} still missing",
+                submitted=total.submitted, still_missing=total.still_missing)
+    return total
+
+
+def wu_full() -> CaptureResult:
+    """
+    Trigger a full Windows Update scan with no HWID filtering.
+    Captures ALL packages WU installs and submits them to the repo.
+    Use to harvest WU drivers onto the repo regardless of current repo state.
+    No scan or resolve needed.
+    """
+    cfg = get()
+    events.start("windows_update",
+                 "[wu-full] Triggering full Windows Update scan (no HWID filtering)…")
+
+    if not _trigger_wu():
+        return CaptureResult()
+
+    wu_start = datetime.now(timezone.utc)
+    total    = _poll_and_submit_incremental(wu_start, cfg)
+    total.wu_triggered = True
+    _post_wu_session(cfg, wu_start, [], total.submitted)
+    events.done("windows_update",
+                f"[wu-full] Done — {total.submitted} submitted  "
+                f"{total.skipped} skipped  {total.failed} failed",
+                submitted=total.submitted, skipped=total.skipped, failed=total.failed)
+    return total
+
+
+# ══ WU HWID filtering ═════════════════════════════════════════════════════════
+
+def _filter_wu_hwids(
+    not_found: list[str],
+    scan_result: ScanResult,
+    cfg: dict,
+) -> list[str]:
+    """
+    Return only the not_found HWIDs that genuinely need Windows Update:
+      - Device has no driver installed locally (has_driver=False)
+      - Device class is not in wu_ignore_classes config
+      - If HWID is not in scan at all → include (unknown device, safer to try)
+    """
+    ignore_classes: set[str] = {
+        c.lower() for c in cfg.get("wu_ignore_classes", [])
+    }
+
+    hwid_to_device: dict[str, Device] = {}
+    for dev in scan_result.devices:
+        for hwid in dev.hwids:
+            hwid_to_device[hwid.upper()] = dev
+
+    result = []
+    for hwid in not_found:
+        dev = hwid_to_device.get(hwid.upper())
+        if dev is None:
+            result.append(hwid)
+            continue
+        if dev.has_driver:
+            continue
+        if ignore_classes and dev.class_name.lower() in ignore_classes:
+            continue
+        result.append(hwid)
+
+    return result
+
+
+# ══ Package enumeration ════════════════════════════════════════════════════════
+
+def _enumerate_all_packages(cfg: dict) -> list[dict]:
+    try:
+        r = subprocess.run(
+            ["pnputil", "/enum-drivers"],
+            capture_output=True, text=True,
+            timeout=cfg.get("timeout_short", 30),
+        )
+    except FileNotFoundError:
+        events.error("capture", "[capture] pnputil not found — must run on Windows",
+                     error="pnputil_not_found")
+        return []
+    except Exception as e:
+        events.error("capture", f"[capture] pnputil failed: {e}", error=str(e))
+        return []
+    return _parse_enum_drivers(r.stdout)
+
+
+def _parse_enum_drivers(output: str) -> list[dict]:
+    packages: list[dict] = []
+    current:  dict       = {}
+    for line in output.splitlines():
+        stripped = line.strip()
+        if not stripped:
+            if current.get("published_name"):
+                packages.append(dict(current))
+            current = {}
+        elif ":" in stripped:
+            raw_key, _, val = stripped.partition(":")
+            mapped = _PNPUTIL_KEYS.get(raw_key.strip().lower())
+            if mapped:
+                current[mapped] = val.strip()
+    if current.get("published_name"):
+        packages.append(current)
+    return packages
+
+
+def _enumerate_packages_since(cutoff: datetime, cfg: dict) -> list[dict]:
+    """Return packages whose C:\\Windows\\INF mtime >= cutoff."""
+    result = []
+    for pkg in _enumerate_all_packages(cfg):
+        inf_path = _WINDOWS_INF / pkg["published_name"]
+        if not inf_path.exists():
+            continue
+        mtime = datetime.fromtimestamp(inf_path.stat().st_mtime, tz=timezone.utc)
+        if mtime >= cutoff:
+            result.append(pkg)
+    return result
+
+
+def _find_installed_packages_for_hwids(
+    not_found_hwids: set[str],
+    scan_result: ScanResult,
+    cfg: dict,
+) -> list[dict]:
+    locally_installed = {
+        hwid for hwid in not_found_hwids
+        if hwid in scan_result.installed_versions
+    }
+    if not locally_installed:
+        return []
+
+    hwid_to_published = _map_hwids_to_published(locally_installed, cfg)
+    all_packages = {
+        p["published_name"]: p
+        for p in _enumerate_all_packages(cfg)
+        if p.get("published_name")
+    }
+    return [
+        all_packages[pub]
+        for pub in set(hwid_to_published.values())
+        if pub in all_packages
+    ]
+
+
+def _map_hwids_to_published(hwids: set[str], cfg: dict) -> dict[str, str]:
+    mapping: dict[str, str] = {}
+    try:
+        r = subprocess.run(
+            ["pnputil", "/enum-devices", "/drivers"],
+            capture_output=True, text=True,
+            timeout=cfg.get("timeout_short", 30),
+        )
+    except Exception:
+        return mapping
+
+    state: dict = {"published": "", "hwids": [], "in_hwids": False}
+    for line in r.stdout.splitlines():
+        stripped = line.strip()
+        if not stripped:
+            _commit_hwid_block(state, hwids, mapping)
+        else:
+            _parse_hwid_line(line, stripped, state)
+
+    return mapping
+
+
+def _commit_hwid_block(state: dict, hwids: set[str], mapping: dict) -> None:
+    """Flush accumulated HWIDs for the current device into mapping."""
+    if state["published"]:
+        for h in state["hwids"]:
+            if h in hwids:
+                mapping[h] = state["published"]
+    state["published"] = ""
+    state["hwids"]     = []
+    state["in_hwids"]  = False
+
+
+def _parse_hwid_line(line: str, stripped: str, state: dict) -> None:
+    """Update parser state from one non-empty pnputil /enum-devices line."""
+    low = stripped.lower()
+    if low.startswith("published name") and ":" in stripped:
+        state["published"] = stripped.split(":", 1)[1].strip()
+        state["in_hwids"]  = False
+    elif low.startswith(("hardware ids", "compatible ids")):
+        state["in_hwids"] = True
+    elif state["in_hwids"] and "\\" in stripped and not stripped.endswith(":"):
+        state["hwids"].append(stripped.upper())
+    elif ":" in stripped and not line.startswith(" "):
+        state["in_hwids"] = False
+
+
+# ══ Windows Update ═════════════════════════════════════════════════════════════
+
+def _trigger_wu() -> bool:
+    for cmd in (["usoclient", "StartScan"], ["wuauclt", "/detectnow"]):
+        try:
+            subprocess.run(cmd, capture_output=True, timeout=30)
+            events.ok("windows_update", "[wu-update] Windows Update scan triggered")
+            return True
+        except Exception:
+            continue
+    events.error("windows_update", "[wu-update] Could not trigger Windows Update",
+                 error="trigger_failed")
+    return False
+
+
+def _poll_and_submit_incremental(wu_start: datetime, cfg: dict) -> CaptureResult:
+    """
+    Poll every wu_poll_interval_seconds for new DriverStore packages.
+    Submit each new batch immediately. Stop after wu_wait_minutes total
+    OR after 2 consecutive polls with no new packages.
+    """
+    wait_min      = cfg.get("wu_wait_minutes", 10)
+    poll_secs     = cfg.get("wu_poll_interval_seconds", 30)
+    deadline      = wu_start + timedelta(minutes=wait_min)
+    seen:set[str] = set()
+    total         = CaptureResult()
+    empty_polls   = 0
+
+    events.progress("windows_update",
+                    f"[wu-update] Polling every {poll_secs}s, max {wait_min} min…")
+
+    while datetime.now(timezone.utc) < deadline:
+        all_new   = _enumerate_packages_since(wu_start, cfg)
+        new_batch = [
+            p for p in all_new
+            if p.get("published_name", "") not in seen
+        ]
+
+        if new_batch:
+            empty_polls = 0
+            for p in new_batch:
+                seen.add(p.get("published_name", ""))
+            events.progress("windows_update",
+                            f"[wu-update] {len(new_batch)} new package(s) — submitting…",
+                            current=len(new_batch))
+            total.merge(_export_and_submit(new_batch, cfg))
+        else:
+            empty_polls += 1
+            if empty_polls >= 2:
+                events.progress("windows_update",
+                                "[wu-update] 2 consecutive empty polls — stopping")
+                break
+            remaining = int((deadline - datetime.now(timezone.utc)).total_seconds())
+            events.progress("windows_update",
+                            f"[wu-update]   no new drivers — {remaining}s remaining…")
+
+        time.sleep(poll_secs)
+
+    return total
+
+
+# ══ Export pool → Upload pool pipeline ════════════════════════════════════════
+
+def _export_and_submit(packages: list[dict], cfg: dict) -> CaptureResult:
+    """
+    Two-stage pipeline:
+      export pool (parallel_exports) → pnputil /export-driver + pack
+      upload pool (parallel_uploads) → POST /submit
+
+    Deduplicates by published_name before starting.
+    """
+    parallel_exports = cfg.get("parallel_exports", 4)
+    parallel_uploads = cfg.get("parallel_uploads", 8)
+
+    # Dedup by published_name
+    seen:     set[str]   = set()
+    unique:   list[dict] = []
+    for pkg in packages:
+        name = pkg.get("published_name", "")
+        if name and name not in seen:
+            seen.add(name)
+            unique.append(pkg)
+
+    if not unique:
+        return CaptureResult()
+
+    counts = {"submitted": 0, "skipped": 0, "failed": 0}
+    lock   = threading.Lock()
+    upload_q: queue.Queue = queue.Queue(maxsize=parallel_uploads * 2)
+
+    # Start upload workers
+    upload_threads = [
+        threading.Thread(target=_upload_worker, args=(upload_q, counts, lock, cfg),
+                         daemon=True)
+        for _ in range(parallel_uploads)
+    ]
+    for t in upload_threads:
+        t.start()
+
+    # Run export pool — each worker exports + packs then puts payload into upload_q
+    with ThreadPoolExecutor(max_workers=parallel_exports) as ex:
+        futures = [ex.submit(_export_one, pkg, upload_q, counts, lock, cfg)
+                   for pkg in unique]
+        for f in futures:
+            try:
+                f.result()
+            except Exception as e:
+                events.error("dump", f"  [error] export: {e}", error=str(e))
+                with lock:
+                    counts["failed"] += 1
+
+    # Signal upload workers to stop
+    for _ in upload_threads:
+        upload_q.put(None)
+    upload_q.join()
+
+    return CaptureResult(
+        submitted=counts["submitted"],
+        skipped=counts["skipped"],
+        failed=counts["failed"],
+    )
+
+
+def _export_one(pkg: dict, upload_q: queue.Queue,
+                counts: dict, lock: threading.Lock, cfg: dict) -> None:
+    """Thin wrapper: export + pack via _do_export, then enqueue for upload."""
+    published  = pkg["published_name"]
+    class_name = pkg.get("class_name", "").lower()
+    events.progress("dump",
+                    f"  [export] {published}  ({pkg.get('provider_name', '?')}  {class_name})",
+                    published_name=published, provider=pkg.get("provider_name", "?"),
+                    class_name=class_name)
+
+    payload = _do_export(pkg, cfg)
+    if payload is None:
+        with lock:
+            counts["failed"] += 1
+        return
+    upload_q.put(payload)
+
+
+def _do_export(pkg: dict, cfg: dict) -> dict | None:
+    """
+    Run pnputil /export-driver, pack the files, return an upload payload dict.
+    TemporaryDirectory cleans up after packing completes — no double-export.
+    Returns None on any error.
+    """
+    published  = pkg["published_name"]
+    original   = pkg.get("original_name", published)
+    class_name = pkg.get("class_name", "").lower()
+    exts       = _pick_extensions(class_name, cfg)
+    timeout    = cfg.get("timeout_long", 120)
+
+    with tempfile.TemporaryDirectory() as tmp:
+        dest = Path(tmp) / "pkg"
+        dest.mkdir()
+
+        try:
+            r = subprocess.run(
+                ["pnputil", "/export-driver", published, str(dest)],
+                capture_output=True, text=True, timeout=timeout,
+            )
+        except subprocess.TimeoutExpired:
+            events.warn("dump", f"  [warn] export timed out: {published}",
+                        published_name=published, error="timeout")
+            return None
+        except Exception as e:
+            events.warn("dump", f"  [warn] export error {published}: {e}",
+                        published_name=published, error=str(e))
+            return None
+
+        if r.returncode != 0:
+            events.warn("dump", f"  [warn] export failed {published}: {r.stderr[:120]}",
+                        published_name=published, error=r.stderr[:120])
+            return None
+
+        try:
+            archive_bytes, archive_name, encoding = pack_driver_archive(
+                dest, Path(original).stem, exts
+            )
+        except Exception as e:
+            events.warn("dump", f"  [warn] pack error {published}: {e}",
+                        published_name=published, error=str(e))
+            return None
+
+    return {
+        "published":      published,
+        "original":       original,
+        "provider":       pkg.get("provider_name", "Unknown"),
+        "class_name":     class_name,
+        "driver_version": pkg.get("driver_version", ""),
+        "archive_bytes":  archive_bytes,
+        "archive_name":   archive_name,
+        "encoding":       encoding,
+    }
+
+
+def _upload_worker(upload_q: queue.Queue, counts: dict,
+                   lock: threading.Lock, cfg: dict) -> None:
+    """Upload worker: pulls payloads from queue and POSTs to /submit."""
+    repo_url = cfg["local_repo_url"]
+    node_key = cfg.get("node_key", "")
+    timeout  = cfg.get("timeout_medium", 60)
+
+    while True:
+        item = upload_q.get()
+        if item is None:
+            upload_q.task_done()
+            return
+        try:
+            resp = http_upload_package(
+                url=f"{repo_url}/submit/upload",
+                node_key=node_key,
+                metadata={
+                    "published_name": item["published"],
+                    "original_name":  item["original"],
+                    "provider":       item["provider"],
+                    "driver_version": item["driver_version"],
+                    "class_name":     item["class_name"],
+                    "source_type":    "captured_live",
+                    **_machine_meta(),
+                },
+                archive_bytes=item["archive_bytes"],
+                archive_name=item["archive_name"],
+                content_encoding=item["encoding"],
+                timeout=timeout,
+            )
+            # The client is done the instant the local repo acks. It does NOT
+            # interpret upstream queue/scan state — "accepted/queued/buffered"
+            # means handed off successfully; only explicit dedup statuses skip.
+            status = resp.get("status")
+            if status in ("already_approved", "already_present", "already_queued"):
+                key = "skipped"
+            else:  # submitted / accepted / queued / buffered → handed off
+                key = "submitted"
+            with lock:
+                counts[key] += 1
+        except BadRequestError as e:
+            events.warn("upload",
+                        f"  [skip] {item.get('published', '?')} rejected by local_repo: {e}",
+                        published_name=item.get("published", "?"), error=str(e))
+            with lock:
+                counts["failed"] += 1
+        except Exception as e:
+            events.error("upload", f"  [error] upload {item.get('published', '?')}: {e}",
+                         published_name=item.get("published", "?"), error=str(e))
+            with lock:
+                counts["failed"] += 1
+        finally:
+            upload_q.task_done()
+
+
+# ══ Extension selection ════════════════════════════════════════════════════════
+
+def _pick_extensions(class_name: str, cfg: dict) -> list[str]:
+    overrides: dict[str, list[str]] = cfg.get("dump_extensions_overrides", {})
+    for key, exts in overrides.items():
+        if key.lower() in class_name:
+            return exts
+    return cfg.get("dump_extensions", [".inf", ".sys", ".cat"])