herdr-session-switch/tests/test_proxy_env.py

from __future__ import annotations

import base64
import os
from pathlib import Path
import shlex
import subprocess
import sys
import tempfile
import unittest
from unittest.mock import patch

from herdr_hub.cli import Endpoint
from herdr_hub.control import HerdrControl
from herdr_hub.proxy_env import configure_application_proxy, configure_proxy, proxy_assignments, running_proxy_paths


class RemoteProxyTests(unittest.TestCase):
    def test_running_mihomo_custom_config_is_discovered(self):
        with tempfile.TemporaryDirectory() as directory:
            root = Path(directory)
            core = root / "1277"
            core.mkdir()
            (core / "cmdline").write_bytes(
                b"/opt/mihomo-linux-arm64\0-d\0/home/orangepi/repo/clash\0-f\0/home/orangepi/repo/clash/config-systemd.yaml\0"
            )
            unrelated = root / "1278"
            unrelated.mkdir()
            (unrelated / "cmdline").write_bytes(b"python3\0-c\0mihomo\0-f\0/irrelevant.yaml\0")
            self.assertEqual(running_proxy_paths(root), [Path("/home/orangepi/repo/clash/config-systemd.yaml")])

    def test_running_clash_listener_used_when_shell_has_no_proxy(self):
        with tempfile.TemporaryDirectory() as directory:
            path = Path(directory) / "clash-verge.yaml"
            path.write_text("mixed-port: 7897\nbind-address: '*'\nallow-lan: false\n")
            environment = {"no_proxy": "localhost,internal"}
            with patch("herdr_hub.proxy_env.socket.create_connection") as connect:
                source = configure_application_proxy(environment, [path])
            self.assertEqual(source, str(path))
            self.assertEqual(environment["https_proxy"], "http://127.0.0.1:7897")
            self.assertEqual(environment["no_proxy"], "localhost,internal")
            connect.assert_called_once_with(("127.0.0.1", 7897), timeout=1)

    def test_stopped_application_and_authenticated_proxy_not_adopted(self):
        with tempfile.TemporaryDirectory() as directory:
            path = Path(directory) / "config.yaml"
            for contents in ("mixed-port: 7897\n", "mixed-port: 7897\nauthentication:\n  - user:secret\n"):
                path.write_text(contents)
                environment = {}
                with patch("herdr_hub.proxy_env.socket.create_connection", side_effect=ConnectionRefusedError):
                    self.assertIsNone(configure_application_proxy(environment, [path]))
                self.assertEqual(environment, {})

    def test_application_detection_respects_existing_proxy_and_disable(self):
        for value in ("", "http://configured:8080"):
            environment = {"HTTPS_PROXY": value}
            with patch("herdr_hub.proxy_env.socket.create_connection") as connect:
                self.assertIsNone(configure_application_proxy(environment, [Path("/unused")]))
                connect.assert_not_called()
            self.assertEqual(environment, {"HTTPS_PROXY": value})

    def test_literal_exports_and_uppercase_are_read_without_shell_evaluation(self):
        values = proxy_assignments('''
export https_proxy="http://remote-user:secret@127.0.0.1:7890" # local to remote
HTTP_PROXY=http://remote-host:8080
export no_proxy='localhost,127.0.0.1'
export all_proxy=$(dangerous_command)
export all_proxy=http://other:8080; dangerous_command
    export all_proxy=http://inside-function:8080
''')
        self.assertEqual(values, {
            "https_proxy": "http://remote-user:secret@127.0.0.1:7890",
            "http_proxy": "http://remote-host:8080",
            "no_proxy": "localhost,127.0.0.1",
        })

    def test_remote_environment_wins_including_explicit_disable(self):
        with tempfile.TemporaryDirectory() as directory:
            path = Path(directory) / ".bashrc"
            path.write_text("export https_proxy=http://file:7890\nexport all_proxy=http://fallback:7890\n")
            environment = {"HTTPS_PROXY": "http://remote-env:7890", "all_proxy": ""}
            configure_proxy(environment, [path])
        self.assertEqual(environment["https_proxy"], "http://remote-env:7890")
        self.assertEqual(environment["all_proxy"], "")

    def test_wrapper_loads_remote_bashrc_and_never_embeds_local_proxy(self):
        with tempfile.TemporaryDirectory() as directory:
            root = Path(directory)
            marker = root / "should-not-exist"
            remote_proxy = "http://remote-user:remote-secret@127.0.0.1:7891"
            (root / ".bashrc").write_text(
                "case $- in *i*) ;; *) return;; esac\n"
                + f"export https_proxy={shlex.quote(remote_proxy)}\n"
                + f"touch {shlex.quote(str(marker))}\n"
            )
            controller = HerdrControl(Endpoint.remote("host"), "unused")
            controller.platform = "posix"
            child = shlex.join([sys.executable, "-c", "import os; assert os.environ['https_proxy'] == " + repr(remote_proxy)])
            with patch.dict(os.environ, {"https_proxy": "http://local-secret@127.0.0.1:9999"}):
                command = controller._update_proxy(["ssh", "host", child], None)
            self.assertEqual(command[:-1], ["ssh", "host"])
            self.assertNotIn("local-secret", command[-1])
            environment = {k: v for k, v in os.environ.items() if k.lower() not in {
                "http_proxy", "https_proxy", "all_proxy", "no_proxy",
            }}
            environment.update(HOME=directory, SHELL="/bin/bash")
            result = subprocess.run(["sh", "-c", command[-1]], env=environment, capture_output=True, text=True, timeout=10)
            self.assertEqual(result.returncode, 0, result.stderr)
            self.assertIn("使用远端自身代理", result.stdout)
            self.assertNotIn("remote-secret", result.stdout + result.stderr)
            self.assertFalse(marker.exists())

    def test_windows_loads_only_remote_environment_scopes(self):
        controller = HerdrControl(Endpoint.remote("windows-host"), "unused")
        controller.platform = "windows"
        with patch.dict(os.environ, {"https_proxy": "http://local-secret@127.0.0.1:9999"}):
            command = controller._update_proxy(["ssh", "windows-host", "unused"], None)
        script = base64.b64decode(command[-1].split()[-1]).decode("utf-16le")
        for scope in ("Process", "User", "Machine"):
            self.assertIn(f"GetEnvironmentVariable($key,'{scope}')", script)
        self.assertNotIn("local-secret", script)
        self.assertNotIn("-R", command)


if __name__ == "__main__":
    unittest.main()