#!/usr/bin/env python3
"""Reproduce "Missing dependencies for SOCKS support" in Requests, pip and HTTPX on loopback.

Every case gets a fresh virtual environment, an empty working directory, an empty HOME and an
environment that holds only PATH, HOME, LANG and the proxy variables the case names. Proxy URLs
point at 127.0.0.1 only. The only outbound traffic is pip talking to PyPI.

    python3 run_lab.py /path/to/python3.13 [/path/to/another/python ...]

Results are written to results.json next to this script.
"""
import datetime
import json
import os
import pathlib
import platform
import shutil
import subprocess
import sys
import time

LAB = pathlib.Path(__file__).resolve().parent
WORK = pathlib.Path(os.environ.get("LAB_WORK", "/var/tmp/socks-lab"))
TARGET = "http://127.0.0.1:18080/"
SOCKS = "socks5://127.0.0.1:1080"
REQUESTS = "requests==2.34.2"
REQUESTS_SOCKS = "requests[socks]==2.34.2"
HTTPX = "httpx==0.28.1"
HTTPX_SOCKS = "httpx[socks]==0.28.1"

REQ_CODE = f'import requests; print("status", requests.get("{TARGET}", proxies={{"http": "{SOCKS}"}}, timeout=10).status_code)'
REQ_ARG = f"python -c '{REQ_CODE}'"
REQ_ARG_CLOSED = REQ_ARG.replace(":1080", ":1081")
REQ_ARG_HTTPS = f"""python -c 'import requests; print("status", requests.get("https://127.0.0.1:18443/", proxies={{"https": "{SOCKS}"}}, timeout=10).status_code)'"""
REQ_ENV = f"""python -c 'import requests; print("status", requests.get("{TARGET}", timeout=10).status_code)'"""
REQ_NOENV = f"""python -c 'import requests; session = requests.Session(); session.trust_env = False; print("status", session.get("{TARGET}", timeout=10).status_code)'"""
HX_ARG = f"""python -c 'import httpx; print("status", httpx.get("{TARGET}", proxy="{SOCKS}", timeout=10).status_code)'"""
HX_ENV = f"""python -c 'import httpx; print("status", httpx.get("{TARGET}", timeout=10).status_code)'"""
HX_CLIENT = """python -c 'import httpx; client = httpx.Client(); print("client created")'"""
HX_NOENV = f"""python -c 'import httpx; print("status", httpx.get("{TARGET}", trust_env=False, timeout=10).status_code)'"""
HX_ASYNC = f"""python -c '
import asyncio, httpx

async def main():
    async with httpx.AsyncClient(proxy="{SOCKS}", timeout=10) as client:
        print("status", (await client.get("{TARGET}")).status_code)

asyncio.run(main())
'"""
CHECK = 'python -c "import socks, sys; print(sys.executable, socks.__file__)"'
CHECK_PATH = 'python -c "import socks; print(socks.__path__)"'
CHECK_HX = 'python -c "import socksio, sys; print(sys.executable, socksio.__file__)"'
PIP_FIX = 'python -m pip install "requests[socks]"'
ENV_GREP = "env | grep -i _proxy"

UP = {"ALL_PROXY": SOCKS}
LOW = {"all_proxy": SOCKS}
BOTH = {"ALL_PROXY": SOCKS, "all_proxy": SOCKS}


def step(sh, env=None):
    return {"sh": sh, "env": env or {}}


def shadow(case_id, setup, files, run=None, script_dir=False):
    run = run or REQ_ARG
    return {"id": case_id, "setup": setup, "files": files, "script_dir": script_dir, "steps": [step(CHECK), step(run)]}


CASES = [
    # Requests without PySocks
    {"id": "requests-proxies-arg", "setup": [REQUESTS], "steps": [step(ENV_GREP), step(REQ_ARG)]},
    {"id": "requests-proxies-arg-https", "setup": [REQUESTS], "steps": [step(REQ_ARG_HTTPS)]},
    {"id": "requests-env-upper", "setup": [REQUESTS], "steps": [step(ENV_GREP, UP), step(REQ_ENV, UP)]},
    {"id": "requests-env-lower", "setup": [REQUESTS], "steps": [step(ENV_GREP, LOW), step(REQ_ENV, LOW)]},
    {"id": "requests-env-http-proxy", "setup": [REQUESTS], "steps": [step(REQ_ENV, {"HTTP_PROXY": SOCKS})]},
    {"id": "requests-env-upper-socks5h", "setup": [REQUESTS], "steps": [step(REQ_ENV, {"ALL_PROXY": "socks5h://127.0.0.1:1080"})]},
    {"id": "requests-env-upper-empty", "setup": [REQUESTS], "steps": [step(REQ_ENV, {"ALL_PROXY": ""})]},
    {"id": "requests-env-lower-empty", "setup": [REQUESTS], "steps": [step(REQ_ENV, {"all_proxy": ""})]},
    {"id": "requests-env-upper-set-lower-empty", "setup": [REQUESTS], "steps": [step(REQ_ENV, {"ALL_PROXY": SOCKS, "all_proxy": ""})]},
    {"id": "requests-env-lower-set-upper-empty", "setup": [REQUESTS], "steps": [step(REQ_ENV, {"all_proxy": SOCKS, "ALL_PROXY": ""})]},
    {"id": "requests-env-upper-with-http-proxy", "setup": [REQUESTS], "steps": [step(REQ_ENV, {"ALL_PROXY": SOCKS, "HTTP_PROXY": "http://127.0.0.1:18081"})]},
    {"id": "requests-env-upper-with-no-proxy", "setup": [REQUESTS], "steps": [step(REQ_ENV, {"ALL_PROXY": SOCKS, "NO_PROXY": "127.0.0.1"})]},
    {"id": "requests-env-upper-trust-env-false", "setup": [REQUESTS], "steps": [step(REQ_NOENV, UP)]},
    {"id": "requests-socksio-only", "setup": [REQUESTS, "socksio==1.0.0"], "steps": [step(REQ_ARG)]},
    # Requests with PySocks
    {"id": "requests-socks-extra", "setup": [REQUESTS_SOCKS], "steps": [step(CHECK), step(REQ_ARG), step(REQ_ENV, UP), step(REQ_ENV, LOW)]},
    {"id": "requests-socks-extra-closed-port", "setup": [REQUESTS_SOCKS], "steps": [step(REQ_ARG_CLOSED)]},
    {"id": "requests-install-fix", "setup": [REQUESTS], "steps": [step(CHECK), step(REQ_ARG), step(PIP_FIX), step(CHECK), step(REQ_ARG)]},
    {"id": "requests-install-pysocks-alone", "setup": [REQUESTS], "steps": [step("python -m pip install pysocks"), step(REQ_ARG)]},
    {"id": "requests-wrong-interpreter", "setup": [REQUESTS], "other_setup": [REQUESTS], "steps": [
        step('pip install "requests[socks]"'), step(REQ_ARG), step("pip --version"), step("python -m pip --version"),
        step(CHECK), step(PIP_FIX), step(CHECK), step(REQ_ARG)]},
    # A local file or folder named socks
    shadow("shadow-empty-dir-pysocks", [REQUESTS_SOCKS], {"socks/": None}),
    shadow("shadow-dir-with-module-pysocks", [REQUESTS_SOCKS], {"socks/server.py": ""}),
    shadow("shadow-package-dir-pysocks", [REQUESTS_SOCKS], {"socks/__init__.py": ""}),
    shadow("shadow-file-pysocks", [REQUESTS_SOCKS], {"socks.py": ""}),
    shadow("shadow-file-pysocks-script", [REQUESTS_SOCKS], {"socks.py": "", "app.py": REQ_CODE + "\n"}, run="python app.py"),
    shadow("shadow-package-dir-pysocks-script", [REQUESTS_SOCKS], {"socks/__init__.py": "", "app.py": REQ_CODE + "\n"}, run="python app.py"),
    shadow("shadow-file-pysocks-script-elsewhere", [REQUESTS_SOCKS], {"socks.py": ""}, run="python {scripts}/app.py", script_dir=True),
    shadow("shadow-empty-dir-no-pysocks", [REQUESTS], {"socks/": None}),
    {"id": "shadow-empty-dir-no-pysocks-path", "setup": [REQUESTS], "files": {"socks/": None}, "steps": [step(CHECK_PATH)]},
    shadow("shadow-file-no-pysocks", [REQUESTS], {"socks.py": ""}),
    # pip
    {"id": "pip-env-lower", "setup": [], "steps": [step(PIP_FIX, LOW), step("python -m pip install pysocks", LOW),
                                                   step("python -m pip download pysocks", LOW), step("python -m pip index versions pysocks", LOW)]},
    {"id": "pip-env-upper", "setup": [], "steps": [step(PIP_FIX, UP)]},
    {"id": "pip-env-upper-empty", "setup": [], "steps": [step(PIP_FIX, {"ALL_PROXY": ""})]},
    {"id": "pip-inline-clear", "setup": [], "steps": [
        step(PIP_FIX, BOTH), step('ALL_PROXY= all_proxy= python -m pip install "requests[socks]"', BOTH), step(CHECK, BOTH),
        step("python -m pip download --no-deps --no-cache-dir --dest wheels idna", BOTH), step(REQ_ENV, BOTH)]},
    {"id": "pip-inline-clear-upper-only", "setup": [], "steps": [step('ALL_PROXY= python -m pip install "requests[socks]"', BOTH)]},
    {"id": "pip-inline-clear-lower-only", "setup": [], "steps": [step('all_proxy= python -m pip install "requests[socks]"', BOTH)]},
    {"id": "pip-env-u", "setup": [], "steps": [step('env -u ALL_PROXY -u all_proxy python -m pip install "requests[socks]"', BOTH)]},
    {"id": "pip-socks-with-pysocks", "setup": ["PySocks==1.7.1"], "steps": [
        step("python -m pip --version"), step("python -m pip install --no-cache-dir idna", LOW),
        step("python -m pip install --no-cache-dir idna", {"all_proxy": "socks5h://127.0.0.1:1080"})]},
    {"id": "pip-bundled-socks-with-pysocks", "bundled_pip": True, "setup": ["PySocks==1.7.1"], "steps": [
        step("python -m pip --version"), step("python -m pip install --no-cache-dir idna", LOW)]},
    # HTTPX
    {"id": "httpx-proxy-arg", "setup": [HTTPX], "steps": [step(ENV_GREP), step(HX_ARG)]},
    {"id": "httpx-env-upper", "setup": [HTTPX], "steps": [step(HX_ENV, UP), step(HX_CLIENT, UP)]},
    {"id": "httpx-env-lower", "setup": [HTTPX], "steps": [step(HX_ENV, LOW)]},
    {"id": "httpx-env-upper-empty", "setup": [HTTPX], "steps": [step(HX_ENV, {"ALL_PROXY": ""})]},
    {"id": "httpx-env-upper-trust-env-false", "setup": [HTTPX], "steps": [step(HX_NOENV, UP)]},
    {"id": "httpx-async-proxy-arg", "setup": [HTTPX], "steps": [step(HX_ASYNC)]},
    {"id": "httpx-pysocks-only", "setup": [HTTPX, "PySocks==1.7.1"], "steps": [step(HX_ARG)]},
    {"id": "httpx-socks-extra", "setup": [HTTPX_SOCKS], "steps": [step(CHECK_HX), step(HX_ARG), step(HX_ENV, UP), step(HX_ENV, LOW), step(HX_ASYNC)]},
    {"id": "httpx-install-fix", "setup": [HTTPX], "steps": [step(HX_ARG), step('python -m pip install "httpx[socks]"'), step(CHECK_HX), step(HX_ARG)]},
]

VERSIONS = """
import json, ssl, sys
from importlib import metadata
found = {}
for name in ("pip", "requests", "urllib3", "PySocks", "httpx", "httpcore", "socksio"):
    try:
        found[name] = metadata.version(name)
    except metadata.PackageNotFoundError:
        found[name] = None
print(json.dumps({"python": sys.version.split()[0], "openssl": ssl.OPENSSL_VERSION, "packages": found}))
"""


def run(cmd, env, cwd, timeout=300):
    done = subprocess.run(cmd, env=env, cwd=cwd, capture_output=True, text=True, timeout=timeout)
    return done.returncode, done.stdout, done.stderr


def base_env(path, home):
    return {"PATH": path, "HOME": str(home), "LANG": "C.UTF-8"}


def versions(venv, root):
    code, out, err = run([str(venv / "bin/python"), "-c", VERSIONS], base_env("/usr/bin:/bin", root / "setup-home"), root)
    if code:
        raise RuntimeError(err)
    return json.loads(out)


def make_venv(python, venv, packages, root, bundled_pip=False):
    env = base_env("/usr/bin:/bin", root / "setup-home")
    upgrade = [] if bundled_pip else [[str(venv / "bin/python"), "-m", "pip", "install", "--quiet", "--upgrade", "pip"]]
    for cmd in ([python, "-m", "venv", str(venv)], *upgrade):
        code, out, err = run(cmd, env, root)
        if code:
            raise RuntimeError(f"{cmd}: {out}{err}")
    if packages:
        code, out, err = run([str(venv / "bin/python"), "-m", "pip", "install", "--quiet", *packages], env, root)
        if code:
            raise RuntimeError(f"install {packages}: {out}{err}")


def run_python(python, relay_log):
    label = subprocess.run([python, "-c", "import sys; print(sys.version.split()[0])"], capture_output=True, text=True).stdout.strip()
    base_prefix = subprocess.run([python, "-c", "import sys; print(sys.base_prefix)"], capture_output=True, text=True).stdout.strip()
    root = WORK / f"py{label}"
    shutil.rmtree(root, ignore_errors=True)
    (root / "setup-home").mkdir(parents=True)
    scripts = root / "scripts"
    scripts.mkdir()
    (scripts / "app.py").write_text(REQ_CODE + "\n")

    def scrub(text):
        return text.replace(base_prefix, f"<python-{label}>") if base_prefix not in ("/usr", "") else text

    cases = []
    for case in CASES:
        started = time.time()
        venv, cwd, home = root / "venvs" / case["id"], root / "cwd" / case["id"], root / "home" / case["id"]
        cwd.mkdir(parents=True)
        home.mkdir(parents=True)
        make_venv(python, venv, case["setup"], root, case.get("bundled_pip", False))
        path = f"{venv}/bin:/usr/bin:/bin"
        if "other_setup" in case:
            other, pathbin = root / "venvs" / (case["id"] + "-other"), root / "pathbin" / case["id"]
            make_venv(python, other, case["other_setup"], root)
            pathbin.mkdir(parents=True)
            (pathbin / "pip").symlink_to(other / "bin/pip")
            path = f"{pathbin}:{path}"
        for name, content in case.get("files", {}).items():
            target = cwd / name
            if content is None:
                target.mkdir(parents=True)
            else:
                target.parent.mkdir(parents=True, exist_ok=True)
                target.write_text(content)
        record = {"id": case["id"], "setup": case["setup"], "working_directory": sorted(case.get("files", {})), "before": versions(venv, root), "steps": []}
        if case.get("bundled_pip"):
            record["note"] = "pip is the version bundled with this Python; it is not upgraded"
        if "other_setup" in case:
            record["note"] = "pip on PATH belongs to a second virtual environment; python belongs to the first"
        for item in case["steps"]:
            env = base_env(path, home) | item["env"]
            offset = relay_log.stat().st_size
            sh = item["sh"].replace("{scripts}", str(scripts))
            code, out, err = run(["/bin/bash", "-c", sh], env, cwd)
            time.sleep(0.2)
            with relay_log.open() as handle:
                handle.seek(offset)
                relay = handle.read().splitlines()
            record["steps"].append({"env": item["env"], "command": sh, "exit": code, "stdout": scrub(out), "stderr": scrub(err), "socks_relay_log": relay})
        record["after"] = versions(venv, root)
        for used in (venv, root / "venvs" / (case["id"] + "-other"), home):
            shutil.rmtree(used, ignore_errors=True)  # keep disk use flat: one environment at a time
        cases.append(record)
        last = record["steps"][-1]
        tail = (last["stderr"].strip() or last["stdout"].strip()).splitlines()[-1:] or [""]
        print(f"[{label}] {case['id']} ({time.time() - started:.0f}s) exit={last['exit']} {tail[0][:150]}", flush=True)
    return {"python": label, "cases": cases}


def main():
    pythons = sys.argv[1:] or [sys.executable]
    WORK.mkdir(parents=True, exist_ok=True)
    www = WORK / "www"
    www.mkdir(exist_ok=True)
    (www / "index.html").write_text("ok\n")
    relay_log = WORK / "relay.log"
    with relay_log.open("w") as log:
        relay = subprocess.Popen([sys.executable, str(LAB / "socks5_relay.py"), "1080"], stdout=log, stderr=subprocess.STDOUT)
    target = subprocess.Popen([sys.executable, "-m", "http.server", "18080", "--bind", "127.0.0.1", "--directory", str(www)],
                              stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
    try:
        time.sleep(1)
        runs = [run_python(python, relay_log) for python in pythons]
    finally:
        relay.terminate()
        target.terminate()
        shutil.rmtree(WORK, ignore_errors=True)
    uname = platform.uname()
    result = {
        "recorded": datetime.datetime.now(datetime.timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ"),
        "system": f"{platform.freedesktop_os_release().get('PRETTY_NAME', uname.system)} ({uname.machine})",
        "target": TARGET,
        "socks_relay": "127.0.0.1:1080 (socks5_relay.py, no authentication); 127.0.0.1:1081 and 127.0.0.1:18081 are closed ports",
        "environment": "PATH, HOME and LANG plus the variables listed in each step; fresh virtual environment, empty working directory and empty HOME per case",
        "runs": runs,
    }
    (LAB / "results.json").write_text(json.dumps(result, indent=1, ensure_ascii=False) + "\n")
    print("wrote", LAB / "results.json")


if __name__ == "__main__":
    main()