stack/tests/prisma/test_vpn_exercise.py

"""Exercise prisma/vpn.py — covers up/down/active/sidecar/prefs/verify."""

from __future__ import annotations

import json
from unittest.mock import MagicMock, patch

import pytest

import prisma.vpn as vpn


class TestDoClient:
    @patch.dict("os.environ", {"DIGITAL_OCEAN_PAT": "test-pat"})
    @patch("pydo.Client")
    def test_creates(self, mc_pydo):
        vpn._do_client()
        mc_pydo.assert_called_once()

    @patch.dict("os.environ", {"DIGITAL_OCEAN_PAT": ""}, clear=False)
    def test_raises(self):
        with pytest.raises(RuntimeError):
            vpn._do_client()


class TestUp:
    @patch("prisma.vpn._do_client")
    @patch("prisma.vpn._gen_ssh_key", return_value="ssh-ed25519 AAAA test")
    @patch("prisma.vpn._ensure_account_ssh_key", return_value=42)
    @patch("prisma.vpn._wait_for_active", return_value="1.2.3.4")
    @patch("prisma.vpn._wait_for_tunnel_ready")
    def test_creates_droplet(
        self, mc_tunnel, mc_wait, mc_ssh, mc_gen, mc_client, tmp_path
    ):
        client = mc_client.return_value
        client.droplets.create.return_value = {"droplet": {"id": 999}}

        state = tmp_path / "state"
        state.mkdir()
        droplet_json = state / "droplet.json"
        env_file = state / "env.sh"

        with (
            patch.object(vpn, "_STATE_DIR", state),
            patch.object(vpn, "_DROPLET_JSON", droplet_json),
            patch.object(vpn, "_ENV_FILE", env_file),
            patch.object(vpn, "_SSH_KEY", state / "key"),
        ):
            info = vpn.up(attach_zotero=False)
        assert info["droplet_id"] == 999
        assert info["public_ip"] == "1.2.3.4"

    def test_raises_if_exists(self, tmp_path):
        dj = tmp_path / "droplet.json"
        dj.write_text("{}")
        with patch.object(vpn, "_DROPLET_JSON", dj):
            with pytest.raises(RuntimeError, match="already tracked"):
                vpn.up()


class TestDown:
    @patch("prisma.vpn._tunnel_stop")
    @patch("prisma.vpn.detach_zotero_proxy")
    @patch("prisma.vpn._do_client")
    @patch("prisma.vpn._delete_account_ssh_key")
    def test_destroys(self, mc_del, mc_client, mc_detach, mc_stop, tmp_path):
        dj = tmp_path / "droplet.json"
        dj.write_text(json.dumps({"id": 123, "ssh_key_id": 42}))

        for f in ["key", "key.pub", "known_hosts", "env.sh", "tunnel.pid"]:
            (tmp_path / f).write_text("")

        with (
            patch.object(vpn, "_DROPLET_JSON", dj),
            patch.object(vpn, "_SSH_KEY", tmp_path / "key"),
            patch.object(vpn, "_KNOWN_HOSTS", tmp_path / "known_hosts"),
            patch.object(vpn, "_ENV_FILE", tmp_path / "env.sh"),
            patch.object(vpn, "_TUNNEL_PID", tmp_path / "tunnel.pid"),
        ):
            result = vpn.down()
        assert result["status"] == "destroyed"

    @patch("prisma.vpn._tunnel_stop")
    @patch("prisma.vpn.detach_zotero_proxy")
    def test_noop(self, mc_detach, mc_stop, tmp_path):
        dj = tmp_path / "droplet.json"
        with patch.object(vpn, "_DROPLET_JSON", dj):
            result = vpn.down()
        assert result["status"] == "nothing-to-do"


class TestStatus:
    def test_up(self, tmp_path):
        dj = tmp_path / "droplet.json"
        dj.write_text(json.dumps({"id": 1, "region": "ams3", "public_ip": "1.2.3.4"}))
        with (
            patch.object(vpn, "_DROPLET_JSON", dj),
            patch("prisma.vpn._tunnel_running", return_value=True),
        ):
            result = vpn.status()
        assert result["status"] == "up"
        assert result["tunnel_running"] is True

    def test_down(self, tmp_path):
        dj = tmp_path / "nope.json"
        with patch.object(vpn, "_DROPLET_JSON", dj):
            result = vpn.status()
        assert result["status"] == "down"


class TestTunnelRunning:
    def test_no_pid_file(self, tmp_path):
        with patch.object(vpn, "_TUNNEL_PID", tmp_path / "nope.pid"):
            assert vpn._tunnel_running() is False

    def test_stale_pid(self, tmp_path):
        pf = tmp_path / "tunnel.pid"
        pf.write_text("9999999")
        with (
            patch.object(vpn, "_TUNNEL_PID", pf),
            patch("os.kill", side_effect=ProcessLookupError),
        ):
            assert vpn._tunnel_running() is False

    def test_live_pid(self, tmp_path):
        pf = tmp_path / "tunnel.pid"
        pf.write_text("1234")
        with (
            patch.object(vpn, "_TUNNEL_PID", pf),
            patch("os.kill"),
        ):
            assert vpn._tunnel_running() is True


class TestTunnelStart:
    @patch("prisma.vpn.subprocess.run")
    @patch("prisma.vpn._tunnel_running", return_value=False)
    @patch("prisma.vpn._pid_listening_on", return_value=5678)
    def test_starts(self, mc_pid, mc_running, mc_sub, tmp_path):
        pf = tmp_path / "tunnel.pid"
        with (
            patch.object(vpn, "_TUNNEL_PID", pf),
            patch.object(vpn, "_SSH_KEY", tmp_path / "key"),
            patch.object(vpn, "_KNOWN_HOSTS", tmp_path / "known_hosts"),
        ):
            pid = vpn._tunnel_start("1.2.3.4")
        assert pid == 5678
        assert pf.read_text() == "5678"


class TestTunnelStop:
    def test_kills(self, tmp_path):
        pf = tmp_path / "tunnel.pid"
        pf.write_text("1234")
        with (
            patch.object(vpn, "_TUNNEL_PID", pf),
            patch("os.kill") as mc_kill,
        ):
            vpn._tunnel_stop()
        mc_kill.assert_called_once_with(1234, 15)


class TestPidListeningOn:
    @patch("prisma.vpn.subprocess.check_output")
    def test_found(self, mc_out):
        mc_out.return_value = (
            'LISTEN 0 128 127.0.0.1:1080 *:* users:(("ssh",pid=12345,fd=4))'
        )
        assert vpn._pid_listening_on(1080) == 12345

    @patch("prisma.vpn.subprocess.check_output", side_effect=FileNotFoundError)
    def test_not_found(self, mc_out):
        assert vpn._pid_listening_on(1080) is None


class TestActive:
    @patch("prisma.vpn._tunnel_running", return_value=False)
    @patch("prisma.vpn._tunnel_start", return_value=1234)
    @patch("prisma.vpn._wait_for_local_port_open")
    @patch("prisma.vpn._tunnel_stop")
    def test_opens_and_closes(self, mc_stop, mc_wait, mc_start, mc_running, tmp_path):
        dj = tmp_path / "droplet.json"
        dj.write_text(json.dumps({"id": 1, "public_ip": "1.2.3.4"}))
        with (
            patch.object(vpn, "_DROPLET_JSON", dj),
        ):
            with vpn.active() as url:
                assert "socks5" in url
        mc_stop.assert_called_once()

    @patch("prisma.vpn._tunnel_running", return_value=True)
    @patch("prisma.vpn._tunnel_stop")
    def test_reuses_existing(self, mc_stop, mc_running, tmp_path):
        dj = tmp_path / "droplet.json"
        dj.write_text(json.dumps({"id": 1, "public_ip": "1.2.3.4"}))
        with patch.object(vpn, "_DROPLET_JSON", dj):
            with vpn.active() as url:
                assert url is not None
        mc_stop.assert_not_called()


class TestSidecar:
    @patch("prisma.vpn._stop_sidecar")
    @patch("prisma.vpn.subprocess.run")
    def test_start(self, mc_run, mc_stop, tmp_path):
        ok_result = MagicMock(returncode=0, stdout="up")
        mc_run.side_effect = [MagicMock(returncode=0), ok_result]
        with patch.object(vpn, "_STATE_DIR", tmp_path):
            result = vpn._start_sidecar("1.2.3.4")
        assert result == vpn._TUNNEL_CONTAINER

    @patch("prisma.vpn.subprocess.run")
    def test_stop(self, mc_run):
        vpn._stop_sidecar()
        mc_run.assert_called_once()


class TestZoteroPrefs:
    @patch("prisma.vpn._rewrite_prefs")
    @patch("prisma.vpn._restart_zotero")
    def test_patch_no_file(self, mc_restart, mc_rewrite, tmp_path):
        with patch.object(vpn, "_ZOTERO_PREFS", tmp_path / "nope"):
            vpn._patch_zotero_prefs()
        mc_rewrite.assert_not_called()

    @patch("prisma.vpn._rewrite_prefs")
    @patch("prisma.vpn._restart_zotero")
    @patch("shutil.copy2")
    def test_patch_with_file(self, mc_copy, mc_restart, mc_rewrite, tmp_path):
        prefs = tmp_path / "prefs.js"
        prefs.write_text('user_pref("foo", "bar");')
        backup = tmp_path / "prefs.js.bak"
        with (
            patch.object(vpn, "_ZOTERO_PREFS", prefs),
            patch.object(vpn, "_ZOTERO_PREFS_BACKUP", backup),
        ):
            vpn._patch_zotero_prefs()
        mc_rewrite.assert_called_once()

    @patch("prisma.vpn.subprocess.run")
    def test_rewrite_prefs(self, mc_run):
        mc_run.return_value = MagicMock(stdout='user_pref("existing", 1);\n')
        vpn._rewrite_prefs({"network.proxy.type": 1})
        assert mc_run.call_count == 3

    @patch("shutil.copy2")
    @patch("pathlib.Path.unlink")
    def test_restore_from_backup(self, mc_unlink, mc_copy, tmp_path):
        prefs = tmp_path / "prefs.js"
        prefs.write_text("content")
        backup = tmp_path / "prefs.js.bak"
        backup.write_text("original")
        with (
            patch.object(vpn, "_ZOTERO_PREFS", prefs),
            patch.object(vpn, "_ZOTERO_PREFS_BACKUP", backup),
        ):
            vpn._restore_zotero_prefs()
        mc_copy.assert_called_once()

    @patch("prisma.vpn._rewrite_prefs")
    def test_restore_no_backup(self, mc_rewrite, tmp_path):
        prefs = tmp_path / "prefs.js"
        prefs.write_text("content")
        backup = tmp_path / "prefs.js.bak"
        with (
            patch.object(vpn, "_ZOTERO_PREFS", prefs),
            patch.object(vpn, "_ZOTERO_PREFS_BACKUP", backup),
        ):
            vpn._restore_zotero_prefs()
        mc_rewrite.assert_called_once()


class TestRestartZotero:
    @patch("prisma.vpn.subprocess.run")
    def test_runs(self, mc_run):
        vpn._restart_zotero()
        mc_run.assert_called_once()


class TestVerifyEgress:
    @patch("httpx.Client")
    def test_returns_json(self, mc_cls):
        client = MagicMock()
        client.__enter__ = MagicMock(return_value=client)
        client.__exit__ = MagicMock(return_value=False)
        resp = MagicMock()
        resp.json.return_value = {"ip": "5.6.7.8", "country": "NL"}
        client.get.return_value = resp
        mc_cls.return_value = client

        result = vpn.verify_egress("socks5://localhost:1080")
        assert result["ip"] == "5.6.7.8"


class TestWaitForActive:
    @patch("prisma.vpn.time.sleep")
    def test_succeeds(self, mc_sleep):
        client = MagicMock()
        client.droplets.get.return_value = {
            "droplet": {
                "status": "active",
                "networks": {"v4": [{"type": "public", "ip_address": "9.8.7.6"}]},
            }
        }
        ip = vpn._wait_for_active(client, 123)
        assert ip == "9.8.7.6"

    @patch("prisma.vpn.time.sleep")
    @patch("prisma.vpn.time.time")
    def test_timeout(self, mc_time, mc_sleep):
        mc_time.side_effect = [0, 0, 999]
        client = MagicMock()
        client.droplets.get.return_value = {
            "droplet": {"status": "new", "networks": {"v4": []}}
        }
        with pytest.raises(TimeoutError):
            vpn._wait_for_active(client, 123, timeout=1)


class TestWaitForTunnelReady:
    @patch("prisma.vpn.time.sleep")
    @patch("prisma.vpn.subprocess.run")
    def test_succeeds(self, mc_sub, mc_sleep, tmp_path):
        mc_sub.return_value = MagicMock(returncode=0, stdout="active")
        vpn._wait_for_tunnel_ready("1.2.3.4", tmp_path / "key", timeout=5)

    @patch("prisma.vpn.time.sleep")
    @patch("prisma.vpn.time.time")
    @patch("prisma.vpn.subprocess.run")
    def test_timeout(self, mc_sub, mc_time, mc_sleep, tmp_path):
        mc_time.side_effect = [0, 0, 999]
        mc_sub.return_value = MagicMock(returncode=1, stdout="")
        with pytest.raises(TimeoutError):
            vpn._wait_for_tunnel_ready("1.2.3.4", tmp_path / "key", timeout=1)


class TestWaitForLocalPort:
    @patch("socket.create_connection")
    def test_succeeds(self, mc_conn):
        ctx = MagicMock()
        ctx.__enter__ = MagicMock()
        ctx.__exit__ = MagicMock(return_value=False)
        mc_conn.return_value = ctx
        vpn._wait_for_local_port_open(1080)

    @patch("prisma.vpn.time.sleep")
    @patch("prisma.vpn.time.time")
    @patch("socket.create_connection", side_effect=OSError)
    def test_timeout(self, mc_conn, mc_time, mc_sleep):
        mc_time.side_effect = [0, 0, 999]
        with pytest.raises(TimeoutError):
            vpn._wait_for_local_port_open(1080, timeout=1)


class TestAttachDetach:
    @patch("prisma.vpn._start_sidecar")
    @patch("prisma.vpn._patch_zotero_prefs")
    @patch("prisma.vpn._restart_zotero")
    def test_attach(self, mc_restart, mc_patch, mc_start):
        vpn.attach_zotero_proxy("1.2.3.4")
        mc_start.assert_called_once()
        mc_patch.assert_called_once()

    @patch("prisma.vpn._stop_sidecar")
    @patch("prisma.vpn._restore_zotero_prefs")
    @patch("prisma.vpn._restart_zotero")
    def test_detach(self, mc_restart, mc_restore, mc_stop):
        vpn.detach_zotero_proxy()
        mc_stop.assert_called_once()
        mc_restore.assert_called_once()