crypto_monitor/scripts/sync_trade_policy_env.py

#!/usr/bin/env python3
"""
将账户方向 / 币种白名单 env 写入三所 .env（缺失则追加，已存在则 --set 时覆盖）。

用法（仓库根目录）：
  python scripts/sync_trade_policy_env.py
  python scripts/sync_trade_policy_env.py --dry-run
  python scripts/sync_trade_policy_env.py --apply-account-profiles
  python scripts/sync_trade_policy_env.py --set-direction binance long_only

--apply-account-profiles：币安=仅多，Gate=BTC/ETH 白名单，OKX=默认不限制。
修改后须 pm2 restart 对应实例。
"""
from __future__ import annotations

import argparse
import os
import re

REPO = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))

INSTANCES = (
    "crypto_monitor_binance",
    "crypto_monitor_okx",
    "crypto_monitor_gate",
)

COMMENT_BLOCK = [
    "# 方向限制（默认 false=双向均可；true 时按 TRADE_DIRECTION 限制，修改后须重启）",
    "# TRADE_DIRECTION=long_only | short_only | both（或 多/空/双向）",
    "# 币种白名单（默认 false=全币种可手输；true 时关键位/下单/策略仅下拉选择）",
]

DEFAULTS = {
    "TRADE_DIRECTION_RESTRICT_ENABLED": "false",
    "TRADE_DIRECTION": "both",
    "TRADE_SYMBOL_RESTRICT_ENABLED": "false",
    "TRADE_SYMBOL_WHITELIST": "BTC,ETH",
}

ACCOUNT_PROFILES = {
    "crypto_monitor_binance": {
        "TRADE_DIRECTION_RESTRICT_ENABLED": "true",
        "TRADE_DIRECTION": "long_only",
        "TRADE_SYMBOL_RESTRICT_ENABLED": "false",
        "TRADE_SYMBOL_WHITELIST": "BTC,ETH",
    },
    "crypto_monitor_gate": {
        "TRADE_DIRECTION_RESTRICT_ENABLED": "false",
        "TRADE_DIRECTION": "both",
        "TRADE_SYMBOL_RESTRICT_ENABLED": "true",
        "TRADE_SYMBOL_WHITELIST": "BTC,ETH",
    },
    "crypto_monitor_okx": {
        "TRADE_DIRECTION_RESTRICT_ENABLED": "false",
        "TRADE_DIRECTION": "both",
        "TRADE_SYMBOL_RESTRICT_ENABLED": "false",
        "TRADE_SYMBOL_WHITELIST": "BTC,ETH",
    },
}


def _parse_env(path: str) -> list[str]:
    if not os.path.isfile(path):
        return []
    with open(path, "r", encoding="utf-8", errors="ignore") as f:
        return f.read().replace("\r\n", "\n").replace("\r", "\n").splitlines()


def _env_get(lines: list[str], key: str) -> str | None:
    pat = re.compile(r"^\s*" + re.escape(key) + r"\s*=\s*(.*)\s*$")
    for line in lines:
        m = pat.match(line)
        if m:
            return m.group(1).strip().strip('"').strip("'")
    return None


def _upsert(lines: list[str], key: str, value: str) -> list[str]:
    pat = re.compile(r"^\s*" + re.escape(key) + r"\s*=")
    out: list[str] = []
    replaced = False
    for line in lines:
        if pat.match(line):
            if not replaced:
                out.append(f"{key}={value}")
                replaced = True
            continue
        out.append(line)
    if not replaced:
        if out and out[-1].strip():
            out.append("")
        out.append(f"{key}={value}")
    return out


def _insert_after(lines: list[str], anchor_key: str, insert: list[str]) -> list[str]:
    pat = re.compile(r"^\s*" + re.escape(anchor_key) + r"\s*=")
    for i, line in enumerate(lines):
        if pat.match(line):
            return lines[: i + 1] + insert + lines[i + 1 :]
    if lines and lines[-1].strip():
        return lines + [""] + insert
    return lines + insert


def sync_one(
    dir_name: str,
    values: dict[str, str],
    *,
    dry_run: bool,
    force: bool,
) -> bool:
    path = os.path.join(REPO, dir_name, ".env")
    if not os.path.isfile(path):
        print(f"skip (no .env): {dir_name}")
        return False
    lines = _parse_env(path)
    changed = False
    for key, val in values.items():
        cur = _env_get(lines, key)
        if cur is None:
            if key == "TRADE_DIRECTION_RESTRICT_ENABLED" and _env_get(
                lines, "TRADE_DIRECTION"
            ) is None:
                if COMMENT_BLOCK[0] not in "\n".join(lines):
                    lines = _insert_after(lines, "POSITION_SIZING_MODE", COMMENT_BLOCK)
            lines = _upsert(lines, key, val)
            changed = True
        elif force or cur != val:
            lines = _upsert(lines, key, val)
            changed = True
    if not changed:
        print(f"ok (unchanged): {dir_name}")
        return False
    text = "\n".join(lines).rstrip() + "\n"
    print(f"update: {dir_name}")
    for k, v in values.items():
        print(f"  {k}={v}")
    if not dry_run:
        with open(path, "w", encoding="utf-8", newline="\n") as f:
            f.write(text)
    return True


def main() -> None:
    ap = argparse.ArgumentParser(description="同步三所 trade policy env")
    ap.add_argument("--dry-run", action="store_true")
    ap.add_argument(
        "--apply-account-profiles",
        action="store_true",
        help="币安仅多、Gate BTC/ETH、OKX 默认",
    )
    ap.add_argument(
        "--defaults-only",
        action="store_true",
        help="三所均写入默认（不限制）",
    )
    ap.add_argument("--force", action="store_true", help="覆盖已有值")
    ap.add_argument("--set-direction", nargs=2, metavar=("INSTANCE", "MODE"))
    ap.add_argument("--set-symbol-whitelist", nargs=2, metavar=("INSTANCE", "SYMS"))
    args = ap.parse_args()

    if args.set_direction:
        inst, mode = args.set_direction
        if inst not in INSTANCES:
            raise SystemExit(f"unknown instance: {inst}")
        sync_one(
            inst,
            {
                "TRADE_DIRECTION_RESTRICT_ENABLED": "true",
                "TRADE_DIRECTION": mode,
            },
            dry_run=args.dry_run,
            force=True,
        )
        return

    if args.set_symbol_whitelist:
        inst, syms = args.set_symbol_whitelist
        if inst not in INSTANCES:
            raise SystemExit(f"unknown instance: {inst}")
        sync_one(
            inst,
            {
                "TRADE_SYMBOL_RESTRICT_ENABLED": "true",
                "TRADE_SYMBOL_WHITELIST": syms,
            },
            dry_run=args.dry_run,
            force=True,
        )
        return

    profiles = (
        {k: dict(DEFAULTS) for k in INSTANCES}
        if args.defaults_only
        else dict(ACCOUNT_PROFILES)
        if args.apply_account_profiles
        else {k: dict(DEFAULTS) for k in INSTANCES}
    )

    if not args.apply_account_profiles and not args.defaults_only:
        ap.print_help()
        print("\n提示：部署常用 --apply-account-profiles")
        return

    any_changed = False
    for inst in INSTANCES:
        if sync_one(inst, profiles[inst], dry_run=args.dry_run, force=args.force):
            any_changed = True
    if args.dry_run and any_changed:
        print("(dry-run, 未写入)")


if __name__ == "__main__":
    main()