tor_updater.py

"""Tor bundle management and update workflow for StealthOps."""

from __future__ import annotations

import hashlib
import json
import os
import platform
import re
import shutil
import subprocess
import sys
import tarfile
import tempfile
import time
import zipfile
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Callable, Optional
from urllib.parse import urljoin, urlparse

import requests


@dataclass
class TorUpdateResult:
    checked: bool
    updated: bool
    message: str
    current_version: Optional[str] = None
    latest_version: Optional[str] = None
    error: Optional[str] = None


class TorUpdater:
    """Manages a writable, app-owned Tor runtime and optional update checks."""

    def __init__(
        self,
        app_name: str = "StealthOps",
        manifest_url: str | None = None,
        ttl_hours: int = 24,
        status_callback: Callable[[str], None] | None = None,
    ) -> None:
        self.manifest_url = manifest_url or os.environ.get("STEALTHOPS_TOR_MANIFEST")
        self.ttl_hours = ttl_hours
        self.status_callback = status_callback
        self.managed_root = self._default_managed_root(app_name)
        self.current_dir = self.managed_root / "current"
        self.state_file = self.managed_root / "update_state.json"

    def _emit_status(self, message: str) -> None:
        if not self.status_callback:
            return
        try:
            self.status_callback(message)
        except Exception:
            pass

    @staticmethod
    def _default_managed_root(app_name: str) -> Path:
        if os.name == "nt":
            base = os.environ.get("LOCALAPPDATA", str(Path.home() / "AppData" / "Local"))
            return Path(base) / app_name / "tor"
        xdg_data = os.environ.get("XDG_DATA_HOME")
        base = Path(xdg_data) if xdg_data else Path.home() / ".local" / "share"
        return base / app_name.lower() / "tor"

    @staticmethod
    def _platform_bundle_suffix() -> str:
        """Return the Tor Project bundle platform suffix for the current OS/arch."""
        if os.name == "nt":
            return "windows-x86_64"
        machine = platform.machine().lower()
        arch = "aarch64" if machine in ("aarch64", "arm64") else "x86_64"
        system = platform.system().lower()
        if system == "darwin":
            return f"macos-{arch}"
        return f"linux-{arch}"

    @property
    def managed_tor_exe(self) -> Path:
        discovered = self._find_tor_executable(self.current_dir)
        if discovered:
            return discovered
        return self.current_dir / ("tor.exe" if os.name == "nt" else "tor")

    def _load_state(self) -> dict[str, Any]:
        if not self.state_file.exists():
            return {}
        try:
            return json.loads(self.state_file.read_text(encoding="utf-8"))
        except Exception:
            return {}

    def _save_state(self, state: dict[str, Any]) -> None:
        self.managed_root.mkdir(parents=True, exist_ok=True)
        self.state_file.write_text(json.dumps(state, indent=2), encoding="utf-8")

    def _check_due(self, mode: str) -> bool:
        if mode == "force":
            return True
        if mode == "off":
            return False

        state = self._load_state()
        last_checked = float(state.get("last_checked", 0))
        return (time.time() - last_checked) >= (self.ttl_hours * 3600)

    def _record_checked(self, latest_version: str | None = None) -> None:
        state = self._load_state()
        state["last_checked"] = time.time()
        if latest_version:
            state["last_seen_version"] = latest_version
        self._save_state(state)

    def get_tor_version(self, tor_path: Path) -> str | None:
        if not tor_path.exists():
            return None
        try:
            completed = subprocess.run(
                [str(tor_path), "--version"],
                capture_output=True,
                text=True,
                timeout=8,
                check=False,
            )
            text = (completed.stdout or "") + " " + (completed.stderr or "")
            match = re.search(r"Tor\s+version\s+([0-9][0-9A-Za-z.\-]+)", text)
            return match.group(1) if match else None
        except Exception:
            return None

    def bootstrap_from_bundle(self, bundle_path: Path | None) -> Path | None:
        """Copies bundled Tor into writable managed location on first run."""
        if self.managed_tor_exe.exists():
            return self.managed_tor_exe

        if not bundle_path or not bundle_path.exists():
            return None

        self.managed_root.mkdir(parents=True, exist_ok=True)
        if self.current_dir.exists():
            shutil.rmtree(self.current_dir, ignore_errors=True)

        source_dir = bundle_path if bundle_path.is_dir() else bundle_path.parent
        shutil.copytree(source_dir, self.current_dir)

        return self._find_tor_executable(self.current_dir)

    @staticmethod
    def _find_tor_executable(root: Path) -> Path | None:
        candidates = [p for p in root.rglob("tor.exe")]
        if not candidates:
            candidates = [
                p for p in root.rglob("tor")
                if p.is_file() and p.parent.name != "debug"
            ]
        return candidates[0] if candidates else None

    @staticmethod
    def _version_key(version: str) -> tuple[int, ...]:
        return tuple(int(part) for part in version.split("."))

    def _fetch_official_tor_project_manifest(self) -> dict[str, Any]:
        page_url = "https://www.torproject.org/download/tor/"
        response = requests.get(page_url, timeout=20)
        response.raise_for_status()
        html_text = response.text

        suffix = self._platform_bundle_suffix()
        pattern = (
            r'href=["\']([^"\']*tor-expert-bundle-'
            + re.escape(suffix)
            + r'-([0-9]+(?:\.[0-9]+)+)\.tar\.gz)["\']'
        )
        matches = re.findall(pattern, html_text)
        if not matches:
            raise RuntimeError(
                f"unable to discover tor expert bundle for {suffix} from torproject.org download page"
            )

        candidates: list[tuple[str, str]] = []
        for href, version in matches:
            absolute_url = urljoin(page_url, href)
            candidates.append((version, absolute_url))

        latest_version, download_url = sorted(candidates, key=lambda item: self._version_key(item[0]))[-1]
        filename = Path(urlparse(download_url).path).name
        sums_url = urljoin(download_url, "sha256sums-signed-build.txt")

        sums_resp = requests.get(sums_url, timeout=20)
        sums_resp.raise_for_status()
        expected_sha = None
        for line in sums_resp.text.splitlines():
            match = re.match(r"^([A-Fa-f0-9]{64})\s+\*?(.+)$", line.strip())
            if match and match.group(2).strip() == filename:
                expected_sha = match.group(1).lower()
                break

        if not expected_sha:
            raise RuntimeError(f"could not locate sha256 for {filename} in {sums_url}")

        return {"version": latest_version, "download_url": download_url, "sha256": expected_sha}

    def _fetch_manifest(self) -> dict[str, Any]:
        manifest_source = self.manifest_url
        source_kind = "explicit" if manifest_source else "auto"
        if not manifest_source:
            candidates = [Path.cwd() / "tor-manifest.json"]
            if getattr(sys, "frozen", False):
                candidates.append(Path(sys.executable).resolve().parent / "tor-manifest.json")

            for candidate in candidates:
                if candidate.exists():
                    manifest_source = str(candidate)
                    source_kind = "local_file"
                    break

        if not manifest_source:
            return self._fetch_official_tor_project_manifest()

        if manifest_source.startswith(("http://", "https://")):
            response = requests.get(manifest_source, timeout=20)
            response.raise_for_status()
            manifest = response.json()
        else:
            manifest = json.loads(Path(manifest_source).read_text(encoding="utf-8"))

        # Accept legacy 'windows_url' key from older manifests
        if "download_url" not in manifest and "windows_url" in manifest:
            manifest["download_url"] = manifest["windows_url"]

        required = ["version", "download_url", "sha256"]
        missing = [key for key in required if key not in manifest]
        if missing:
            raise RuntimeError(f"manifest missing keys: {', '.join(missing)}")
        if "example.com" in str(manifest.get("download_url", "")):
            if source_kind == "local_file":
                return self._fetch_official_tor_project_manifest()
            raise RuntimeError("manifest appears to be a template; replace download_url/sha256 with real values")

        return manifest

    def preview_update_source(self) -> dict[str, Any]:
        """Return the current resolved update candidate without downloading."""
        manifest = self._fetch_manifest()
        return {"version": str(manifest["version"]), "download_url": str(manifest["download_url"])}

    @staticmethod
    def _sha256_file(path: Path) -> str:
        digest = hashlib.sha256()
        with path.open("rb") as handle:
            while True:
                chunk = handle.read(1024 * 1024)
                if not chunk:
                    break
                digest.update(chunk)
        return digest.hexdigest().lower()

    @staticmethod
    def _format_mb(byte_count: int) -> str:
        return f"{(byte_count / (1024 * 1024)):.1f}"

    def _download(self, url: str, destination: Path) -> tuple[int, int | None]:
        host = urlparse(url).netloc or "source"
        self._emit_status(f"downloading tor bundle from {host}")
        with requests.get(url, timeout=60, stream=True) as response:
            response.raise_for_status()
            total_header = response.headers.get("Content-Length", "").strip()
            total = int(total_header) if total_header.isdigit() else None
            downloaded = 0
            last_percent_reported = -1
            last_bytes_reported = 0
            with destination.open("wb") as out:
                for chunk in response.iter_content(chunk_size=1024 * 1024):
                    if chunk:
                        out.write(chunk)
                        downloaded += len(chunk)
                        if total:
                            percent = int((downloaded * 100) / total)
                            if percent >= last_percent_reported + 10 or downloaded >= total:
                                self._emit_status(
                                    f"downloading tor bundle: {self._format_mb(downloaded)} / {self._format_mb(total)} MB ({percent}%)"
                                )
                                last_percent_reported = percent
                        else:
                            if downloaded - last_bytes_reported >= 5 * 1024 * 1024:
                                self._emit_status(f"downloading tor bundle: {self._format_mb(downloaded)} MB")
                                last_bytes_reported = downloaded
        self._emit_status(f"download complete: {self._format_mb(downloaded)} MB")
        return downloaded, total

    def _extract_bundle(self, archive_path: Path, target_dir: Path) -> None:
        if target_dir.exists():
            shutil.rmtree(target_dir, ignore_errors=True)
        target_dir.mkdir(parents=True, exist_ok=True)

        suffixes = "".join(archive_path.suffixes).lower()
        if archive_path.suffix.lower() == ".zip":
            with zipfile.ZipFile(archive_path, "r") as zf:
                zf.extractall(target_dir)
            return

        if ".tar" in suffixes:
            with tarfile.open(archive_path, "r:*") as tf:
                tf.extractall(target_dir)
            return

        raise RuntimeError(f"unsupported tor bundle format: {archive_path.name}")

    def _activate_staged(self, staged_dir: Path) -> None:
        previous_dir = self.managed_root / "previous"

        if previous_dir.exists():
            shutil.rmtree(previous_dir, ignore_errors=True)

        if self.current_dir.exists():
            self.current_dir.rename(previous_dir)

        staged_dir.rename(self.current_dir)

    def maybe_update(self, mode: str = "auto") -> TorUpdateResult:
        if mode not in {"auto", "force", "off"}:
            return TorUpdateResult(False, False, f"invalid update mode: {mode}", error="invalid_mode")

        if mode == "off":
            return TorUpdateResult(False, False, "update checks disabled")

        if not self._check_due(mode):
            current = self.get_tor_version(self.managed_tor_exe)
            return TorUpdateResult(False, False, "update check skipped (ttl)", current_version=current)

        try:
            manifest = self._fetch_manifest()
        except Exception as exc:
            self._record_checked()
            return TorUpdateResult(True, False, f"update check failed: {exc}", error=str(exc))

        latest = str(manifest["version"])
        current = self.get_tor_version(self.managed_tor_exe)
        self._record_checked(latest)

        if current and current == latest:
            return TorUpdateResult(True, False, "tor already up to date", current_version=current, latest_version=latest)

        try:
            self.managed_root.mkdir(parents=True, exist_ok=True)
            with tempfile.TemporaryDirectory(prefix="stealthops-tor-") as tmp_dir:
                temp_dir = Path(tmp_dir)
                archive_name = Path(str(manifest["download_url"])).name or "tor_bundle.tar.gz"
                archive_path = temp_dir / archive_name

                downloaded_bytes, _ = self._download(str(manifest["download_url"]), archive_path)

                expected = str(manifest["sha256"]).lower().strip()
                actual = self._sha256_file(archive_path)
                if expected != actual:
                    raise RuntimeError("sha256 verification failed")

                staged = self.managed_root / "staged"
                self._extract_bundle(archive_path, staged)

                staged_exe = self._find_tor_executable(staged)
                if not staged_exe:
                    raise RuntimeError("updated bundle does not contain tor executable")

                if os.name != "nt":
                    staged_exe.chmod(staged_exe.stat().st_mode | 0o111)

                self._activate_staged(staged)

            updated_version = self.get_tor_version(self.managed_tor_exe)
            return TorUpdateResult(
                checked=True,
                updated=True,
                message=f"tor updated successfully (downloaded {self._format_mb(downloaded_bytes)} MB)",
                current_version=updated_version,
                latest_version=latest,
            )
        except Exception as exc:
            return TorUpdateResult(
                checked=True,
                updated=False,
                message=f"tor update failed: {exc}",
                current_version=current,
                latest_version=latest,
                error=str(exc),
            )