mihomo_injecter/app/expander.py

import yaml
from yaml.representer import SafeRepresenter


def _str_representer(dumper: yaml.Dumper, value: str) -> yaml.ScalarNode:
    # Quote strings that start with special YAML characters so Go parser doesn't choke
    if value and value[0] in ('-', ':', '{', '}', '[', ']', ',', '#', '&', '*', '?', '|', '<', '>', '=', '!', '%', '@', '`'):
        return dumper.represent_scalar('tag:yaml.org,2002:str', value, style="'")
    return SafeRepresenter.represent_str(dumper, value)


class _SafeDumper(yaml.SafeDumper):
    pass


_SafeDumper.add_representer(str, _str_representer)


def _dump(data: dict) -> str:
    return yaml.dump(data, Dumper=_SafeDumper, allow_unicode=True, sort_keys=False)

PROXY_FIELDS = {
    "name", "type", "server", "port", "uuid", "password", "cipher",
    "alterId", "tls", "network", "ws-opts", "h2-opts", "http-opts",
    "grpc-opts", "reality-opts", "servername", "fingerprint", "flow",
    "udp", "skip-cert-verify", "sni", "username", "plugin", "plugin-opts",
    "client-fingerprint", "early-data-header-name", "smux",
}

# Types that are not real proxy protocols — Mihomo may include them in provider
# responses but they're not valid standalone proxy entries.
_SKIP_TYPES = {"direct", "dns", "reject", "pass", "compatible", "relay",
               "selector", "fallback", "url-test", "load-balance"}


def filter_proxy(proxy: dict) -> dict:
    result = {k: v for k, v in proxy.items() if k in PROXY_FIELDS}
    if "type" in result and isinstance(result["type"], str):
        result["type"] = result["type"].lower()
    return result


def _is_real_proxy(proxy: dict) -> bool:
    return "name" in proxy and proxy.get("type", "").lower() not in _SKIP_TYPES


def expand_config(base_yaml: str, provider_proxies: dict[str, list[dict]]) -> str:
    cfg = yaml.safe_load(base_yaml) or {}

    all_proxies = list(cfg.get("proxies") or [])
    provider_to_names: dict[str, list[str]] = {}
    all_sub_names: list[str] = []

    for provider_name, proxies in provider_proxies.items():
        filtered = [filter_proxy(p) for p in proxies if _is_real_proxy(p)]
        provider_to_names[provider_name] = [p["name"] for p in filtered]
        all_proxies.extend(filtered)
        all_sub_names.extend(p["name"] for p in filtered)

    cfg["proxies"] = all_proxies

    # Deduplicate subscription proxy names (preserve order)
    seen: set[str] = set()
    unique_sub_names: list[str] = []
    for n in all_sub_names:
        if n not in seen:
            seen.add(n)
            unique_sub_names.append(n)

    cfg["proxy-groups"] = list(cfg.get("proxy-groups") or [])

    sub_names_set = set(unique_sub_names)
    expanded_group_names: list[str] = []

    for group in cfg.get("proxy-groups") or []:
        group.pop("hidden", None)
        if "use" not in group:
            continue
        name = group.get("name", "")
        expanded: list[str] = []
        for pname in group.pop("use"):
            expanded.extend(provider_to_names.get(pname, []))

        original = group.get("proxies") or []
        special = [p for p in original if p not in sub_names_set and p != "DIRECT"]
        group["proxies"] = special + expanded + ["DIRECT"]
        if not group["proxies"]:
            group["proxies"] = ["DIRECT"]

        if name and name != "GLOBAL":
            expanded_group_names.append(name)

    # Add all service groups to GLOBAL so FLClashX can discover them
    if expanded_group_names:
        global_group = next(
            (g for g in cfg.get("proxy-groups") or [] if g.get("name") == "GLOBAL"),
            None,
        )
        if global_group:
            current = global_group.get("proxies") or []
            current_set = set(current)
            to_add = [n for n in expanded_group_names if n not in current_set]
            global_group["proxies"] = current + to_add

    cfg.pop("proxy-providers", None)

    return _dump(cfg)


def inject_providers_for_delivery(base_yaml: str, subscriptions: list) -> str:
    """At delivery time: inject proxy-providers entries and use: for subscriptions absent from base_yaml."""
    if not subscriptions:
        return base_yaml
    cfg = yaml.safe_load(base_yaml) or {}
    db_sub_names = {sub.name for sub in subscriptions}

    # Remove placeholder providers that aren't real DB subscriptions
    if cfg.get("proxy-providers"):
        for name in list(cfg["proxy-providers"].keys()):
            if name not in db_sub_names:
                del cfg["proxy-providers"][name]

    existing_providers = set((cfg.get("proxy-providers") or {}).keys())
    new_entries = {}
    for sub in subscriptions:
        if sub.name not in existing_providers:
            new_entries[sub.name] = {
                "type": "http",
                "url": sub.url,
                "interval": 3600,
                "health-check": {
                    "enable": True,
                    "url": "https://www.gstatic.com/generate_204",
                    "interval": 300,
                },
            }
    if not new_entries:
        return base_yaml
    all_known = existing_providers | set(new_entries.keys())
    providers = cfg.setdefault("proxy-providers", {})
    providers.update(new_entries)
    new_names = list(new_entries.keys())
    _special = {"DIRECT", "REJECT"}
    for group in cfg.get("proxy-groups") or []:
        if group.get("name") in _special:
            continue
        if "use" in group:
            # Replace references to unknown providers with our subscriptions
            patched = []
            for u in group["use"]:
                if u in all_known:
                    patched.append(u)
                else:
                    patched.extend(new_names)
            group["use"] = patched if patched else new_names
        else:
            group["use"] = new_names
    return _dump(cfg)


def build_mihomo_config(all_base_yamls: list[str], secret: str) -> str:
    providers: dict[str, dict] = {}

    for base_yaml in all_base_yamls:
        try:
            cfg = yaml.safe_load(base_yaml) or {}
            for name, provider in (cfg.get("proxy-providers") or {}).items():
                if name not in providers:
                    providers[name] = provider
        except yaml.YAMLError:
            continue

    config: dict = {
        "external-controller": "0.0.0.0:9090",
        "secret": secret,
        "mixed-port": 7890,
        "allow-lan": False,
    }
    if providers:
        config["proxy-providers"] = providers

    return _dump(config)