crypto_monitor/scripts/sync_four_exchange_transfer_env.py

#!/usr/bin/env python3
"""
将每日自动划转相关项写入四所实例 .env（已有值保留，缺失则追加；可选强制改金额/开关）。

用法（仓库根目录）：
  python scripts/sync_four_exchange_transfer_env.py
  python scripts/sync_four_exchange_transfer_env.py --dry-run
  python scripts/sync_four_exchange_transfer_env.py --set-amount 50
  python scripts/sync_four_exchange_transfer_env.py --enable-auto-transfer

不修改 API 密钥与其它自定义项；若 .env 不存在则跳过（请先从 .env.example 复制）。
"""
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",
    "crypto_monitor_gate_bot",
)

COMMENT_BLOCK = (
    "# 自动划转：北京时间 AUTO_TRANSFER_BJ_HOUR 点将 swap 调整至 AUTO_TRANSFER_AMOUNT；"
    "不足 funding→swap、超出 swap→funding；持仓中不划转"
)

DEFAULTS = {
    "AUTO_TRANSFER_ENABLED": "false",
    "AUTO_TRANSFER_FROM": "funding",
    "AUTO_TRANSFER_TO": "swap",
    "TRANSFER_CCY": "USDT",
    "AUTO_TRANSFER_BJ_HOUR": "8",
}

DEFAULT_AMOUNT = "50"

BINANCE_ONLY = {
    "BINANCE_FUNDING_INCLUDE_SPOT": "false",
}


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 = []
    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_before(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] + insert + lines[i:]
    if lines and lines[-1].strip():
        return lines + [""] + insert
    return lines + insert


def _resolve_default_amount(lines: list[str]) -> str:
    amount = _env_get(lines, "AUTO_TRANSFER_AMOUNT")
    if amount is not None:
        return amount
    daily = _env_get(lines, "DAILY_START_CAPITAL")
    if daily is not None:
        return daily
    return DEFAULT_AMOUNT


def _ensure_key(
    lines: list[str],
    key: str,
    value: str,
    *,
    force: bool,
) -> list[str]:
    if force or _env_get(lines, key) is None:
        return _upsert(lines, key, value)
    return lines


def _ensure_transfer_block(
    lines: list[str],
    extra: dict[str, str],
    *,
    force_amount: str | None,
    force_enabled: str | None,
) -> list[str]:
    amount = force_amount if force_amount is not None else _resolve_default_amount(lines)
    had_amount = _env_get(lines, "AUTO_TRANSFER_AMOUNT") is not None

    if not had_amount and COMMENT_BLOCK not in lines:
        lines = _insert_before(
            lines,
            "AUTO_TRANSFER_ENABLED",
            [COMMENT_BLOCK],
        )
        if _env_get(lines, "AUTO_TRANSFER_ENABLED") is None:
            lines = _insert_before(
                lines,
                "BALANCE_REFRESH_SECONDS",
                [COMMENT_BLOCK],
            )

    lines = _ensure_key(
        lines,
        "AUTO_TRANSFER_AMOUNT",
        amount,
        force=force_amount is not None,
    )
    for k, v in DEFAULTS.items():
        if k == "AUTO_TRANSFER_ENABLED" and force_enabled is not None:
            lines = _upsert(lines, k, force_enabled)
        else:
            lines = _ensure_key(lines, k, v, force=False)
    for k, v in extra.items():
        lines = _ensure_key(lines, k, v, force=False)
    return lines


def sync_one(
    dir_name: str,
    dry_run: bool,
    *,
    set_amount: str | None,
    enable_auto: bool | None,
) -> str:
    env_path = os.path.join(REPO, dir_name, ".env")
    if not os.path.isfile(env_path):
        return f"SKIP {dir_name}: 无 .env（请 cp .env.example .env）"
    old_lines = _parse_env(env_path)
    extra = dict(BINANCE_ONLY) if dir_name == "crypto_monitor_binance" else {}
    force_enabled = "true" if enable_auto is True else None
    new_lines = _ensure_transfer_block(
        old_lines,
        extra,
        force_amount=set_amount,
        force_enabled=force_enabled,
    )
    enabled = _env_get(new_lines, "AUTO_TRANSFER_ENABLED") or DEFAULTS["AUTO_TRANSFER_ENABLED"]
    amt = _env_get(new_lines, "AUTO_TRANSFER_AMOUNT") or DEFAULT_AMOUNT
    if new_lines == old_lines:
        return f"OK   {dir_name}: ENABLED={enabled} AMOUNT={amt}"
    if dry_run:
        return f"DRY  {dir_name}: 将更新 ENABLED={enabled} AMOUNT={amt}"
    with open(env_path, "w", encoding="utf-8", newline="\n") as f:
        f.write("\n".join(new_lines))
        if new_lines and new_lines[-1].strip():
            f.write("\n")
    return f"DONE {dir_name}: ENABLED={enabled} AMOUNT={amt}"


def main():
    ap = argparse.ArgumentParser(description="四所 .env 自动划转项同步")
    ap.add_argument("--dry-run", action="store_true")
    ap.add_argument(
        "--set-amount",
        metavar="U",
        help=f"强制四所 AUTO_TRANSFER_AMOUNT（缺省补全默认 {DEFAULT_AMOUNT}）",
    )
    ap.add_argument(
        "--enable-auto-transfer",
        action="store_true",
        help="强制四所 AUTO_TRANSFER_ENABLED=true",
    )
    args = ap.parse_args()
    for name in INSTANCES:
        print(
            sync_one(
                name,
                args.dry_run,
                set_amount=args.set_amount,
                enable_auto=True if args.enable_auto_transfer else None,
            )
        )


if __name__ == "__main__":
    main()