infra/ansible/files/ru-vps-mihomo-harden.py

#!/usr/bin/env python3
from __future__ import annotations

import argparse
import copy
import json
import os
import secrets
import stat
import subprocess
import sys
import tempfile
from pathlib import Path

import yaml


DEFAULT_PROXY_HOST = "127.0.0.1"
DEFAULT_PROXY_PORT = 7891


def _atomic_write_bytes(path: Path, data: bytes, mode: int) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    fd, tmp_name = tempfile.mkstemp(prefix=f".{path.name}.", dir=str(path.parent))
    try:
        with os.fdopen(fd, "wb") as handle:
            handle.write(data)
            handle.flush()
            os.fsync(handle.fileno())
        os.chmod(tmp_name, mode)
        os.replace(tmp_name, path)
    finally:
        if os.path.exists(tmp_name):
            os.unlink(tmp_name)


def _atomic_write_json(path: Path, payload: dict[str, str]) -> None:
    _atomic_write_bytes(path, json.dumps(payload, indent=2, sort_keys=True).encode("utf-8") + b"\n", 0o600)


def _load_json(path: Path) -> dict[str, str]:
    with path.open("r", encoding="utf-8") as handle:
        payload = json.load(handle)
    if not isinstance(payload, dict):
        raise ValueError(f"{path} must contain a JSON object")
    return payload


def _generate_username() -> str:
    return f"mihomo-{secrets.token_hex(8)}"


def _generate_password() -> str:
    return secrets.token_urlsafe(48)


def ensure_credentials(state_path: Path) -> tuple[dict[str, str], bool]:
    if state_path.exists():
        payload = _load_json(state_path)
        username = payload.get("username")
        password = payload.get("password")
        if not username or not password:
            raise ValueError(f"{state_path} is missing username/password")
        return {"username": username, "password": password}, False

    state_path.parent.mkdir(parents=True, exist_ok=True)
    os.chmod(state_path.parent, 0o700)
    credentials = {"username": _generate_username(), "password": _generate_password()}
    _atomic_write_json(state_path, credentials)
    return credentials, True


def load_config(config_path: Path) -> dict:
    with config_path.open("r", encoding="utf-8") as handle:
        payload = yaml.safe_load(handle) or {}
    if not isinstance(payload, dict):
        raise ValueError(f"{config_path} must contain a YAML mapping")
    return payload


def write_config(config_path: Path, payload: dict) -> None:
    serialized = yaml.safe_dump(payload, sort_keys=False, allow_unicode=True)
    mode = 0o640
    if config_path.exists():
        mode = stat.S_IMODE(config_path.stat().st_mode)
    _atomic_write_bytes(config_path, serialized.encode("utf-8"), mode)


def cmd_apply(args: argparse.Namespace) -> int:
    credentials, state_created = ensure_credentials(args.state)
    config = load_config(args.config)
    updated = copy.deepcopy(config)
    updated["authentication"] = [f"{credentials['username']}:{credentials['password']}"]
    updated["allow-lan"] = False
    updated["bind-address"] = "127.0.0.1"
    changed = updated != config
    if changed:
        write_config(args.config, updated)
    print(json.dumps({"changed": changed, "state_created": state_created}, sort_keys=True))
    return 0


def cmd_probe(args: argparse.Namespace) -> int:
    credentials = ensure_credentials(args.state)[0]
    curl_config = "\n".join(
        [
            'silent',
            'show-error',
            f'proxy = "socks5h://{DEFAULT_PROXY_HOST}:{DEFAULT_PROXY_PORT}"',
            f'proxy-user = "{credentials["username"]}:{credentials["password"]}"',
            'connect-timeout = 5',
            'max-time = 20',
            'url = "https://api.telegram.org"',
            'output = "/dev/null"',
            '',
        ]
    )
    result = subprocess.run(
        ["curl", "--config", "-"],
        input=curl_config,
        text=True,
        capture_output=True,
        timeout=25,
        check=False,
    )
    if result.returncode != 0:
        sys.stderr.write(result.stderr)
        return result.returncode
    print(json.dumps({"ok": True}, sort_keys=True))
    return 0


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    subparsers = parser.add_subparsers(dest="command", required=True)

    for name, func in (("apply", cmd_apply), ("probe", cmd_probe)):
        sub = subparsers.add_parser(name)
        sub.add_argument("--config", required=True, type=Path)
        sub.add_argument("--state", required=True, type=Path)
        sub.set_defaults(func=func)

    return parser.parse_args()


def main() -> int:
    args = parse_args()
    return args.func(args)


if __name__ == "__main__":
    raise SystemExit(main())