websurfer/crawler/tor.py

import atexit
import shutil
import time
from typing import Optional
import stem.process
from stem import Signal
from stem.control import Controller

DEFAULT_TOR_SOCKS_PORT = 9050
DEFAULT_TOR_CONTROL_PORT = 9051


class TorManager:
    """Manages the standalone headless Tor daemon process and circuit rotation."""

    def __init__(
        self,
        socks_port: int = DEFAULT_TOR_SOCKS_PORT,
        control_port: int = DEFAULT_TOR_CONTROL_PORT,
        exit_nodes: Optional[str] = None,
        log_callback: Optional[object] = None,
    ):
        self.socks_port = socks_port
        self.control_port = control_port
        self.exit_nodes = exit_nodes
        self.log_callback = log_callback
        self.process: Optional[object] = None
        self._is_external_service = False

    def set_exit_nodes(self, exit_nodes: Optional[str]) -> None:
        """Dynamically updates exit node country filter on the active Tor controller."""
        self.exit_nodes = exit_nodes
        try:
            with Controller.from_port(port=self.control_port) as controller:
                controller.authenticate()
                if self.exit_nodes:
                    controller.set_conf("ExitNodes", self.exit_nodes)
                    print(f"[Tor] Configured ExitNodes on Tor controller: {self.exit_nodes}")
                else:
                    try:
                        controller.reset_conf("ExitNodes")
                    except Exception:
                        pass
                    print("[Tor] Reset ExitNodes filter (Global / Any)")
        except Exception as e:
            print(f"  [!] Failed to set ExitNodes on ControlPort {self.control_port}: {e}")

    def start(self):
        """
        Starts a background Tor process with custom SocksPort and ControlPort using stem,
        or attaches to an already running Tor service if available.
        """
        # 1. Check if a Tor controller is already running on the control port
        try:
            with Controller.from_port(port=self.control_port) as controller:
                controller.authenticate()
                print(f"[Tor] Connected to existing Tor service on ControlPort {self.control_port}.")
                if self.exit_nodes:
                    controller.set_conf("ExitNodes", self.exit_nodes)
                    print(f"[Tor] Applied ExitNodes: {self.exit_nodes}")
                else:
                    try:
                        controller.reset_conf("ExitNodes")
                    except Exception:
                        pass
                self._is_external_service = True
                return None
        except Exception:
            pass

        # 2. Check if 'tor' binary is installed
        if not shutil.which("tor"):
            raise RuntimeError(
                "Tor executable ('tor') was not found on your system PATH.\n"
                "Please install Tor:\n"
                "  - macOS: brew install tor\n"
                "  - Linux: sudo apt install tor\n"
                "  - Windows: download Tor expert bundle from torproject.org and add to PATH"
            )

        print(
            f"[Tor] Launching headless Tor process "
            f"(SocksPort {self.socks_port}, ControlPort {self.control_port})..."
        )
        # Always launch Tor with base config so it can bootstrap directory info and GeoIP cleanly without stalling
        tor_config = {
            "SocksPort": str(self.socks_port),
            "ControlPort": str(self.control_port),
        }

        def _handle_tor_msg(msg: str):
            if "Bootstrapped" in msg:
                clean_msg = f"[Tor] {msg.strip()}"
                print(f"  {clean_msg}")
                if self.log_callback and callable(self.log_callback):
                    try:
                        self.log_callback(clean_msg)
                    except Exception:
                        pass

        self.process = stem.process.launch_tor_with_config(
            config=tor_config,
            take_ownership=False,  # Managed explicitly via atexit/stop to prevent premature shutdowns
            init_msg_handler=_handle_tor_msg,
        )
        self._is_external_service = False

        # Apply ExitNodes once Tor is successfully bootstrapped
        if self.exit_nodes:
            try:
                with Controller.from_port(port=self.control_port) as controller:
                    controller.authenticate()
                    controller.set_conf("ExitNodes", self.exit_nodes)
                    print(f"[Tor] Applied ExitNodes filter: {self.exit_nodes}")
            except Exception as e:
                print(f"  [!] Failed to apply ExitNodes filter on ControlPort {self.control_port}: {e}")

        # Ensure background process is killed when python script exits
        atexit.register(self.stop)
        print("[Tor] Standalone Tor process started successfully!")
        return self.process

    def renew_circuit(self) -> bool:
        """
        Sends NEWNYM signal to Tor ControlPort to switch to a new circuit / IP.
        Waits for the necessary cool-off period required by Tor.
        Returns True if signal was sent successfully, False otherwise.
        """
        try:
            with Controller.from_port(port=self.control_port) as controller:
                controller.authenticate()
                if self.exit_nodes:
                    controller.set_conf("ExitNodes", self.exit_nodes)
                controller.signal(Signal.NEWNYM)
                wait_time = controller.get_newnym_wait()
                print(f"[Tor] Requested new circuit (NEWNYM). Waiting {wait_time:.1f}s for new circuit to activate...")
                if wait_time > 0:
                    time.sleep(wait_time)
                print("[Tor] New circuit is active!")
                return True
        except Exception as e:
            print(f"  [!] Failed to renew Tor circuit via ControlPort {self.control_port}: {e}")
            return False

    def stop(self) -> None:
        """Stops the background Tor process if it was spawned by this manager."""
        if self.process and not self._is_external_service:
            try:
                if self.process.poll() is None:
                    print("[Tor] Stopping headless Tor process...")
                    self.process.kill()
            except Exception:
                pass
            finally:
                self.process = None

    def __enter__(self):
        self.start()
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.stop()


# Backward-compatible functional interfaces
def start_tor_service(
    socks_port: int = DEFAULT_TOR_SOCKS_PORT,
    control_port: int = DEFAULT_TOR_CONTROL_PORT,
):
    manager = TorManager(socks_port=socks_port, control_port=control_port)
    manager.start()
    return manager


def renew_tor_circuit(control_port: int = DEFAULT_TOR_CONTROL_PORT) -> None:
    TorManager(control_port=control_port).renew_circuit()


def stop_tor_service(tor_process) -> None:
    if tor_process and hasattr(tor_process, "kill"):
        try:
            if tor_process.poll() is None:
                tor_process.kill()
        except Exception:
            pass