xiaotutongxue/scripts/dev-login.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
开发辅助：RSA 加密登录并打印 token。

为什么需要它
------------
登录接口的密码字段是 RSA 加密的（防止明文密码进网关日志），
所以用 curl / Postman 直接调会很别扭——得先算出一段密文。
这个脚本把它变成一行命令。

用法
----
    python scripts/dev-login.py                        # 默认 admin / Admin@123
    python scripts/dev-login.py platform Platform@123
    python scripts/dev-login.py admin Admin@123 --base http://localhost:8080

依赖：cryptography  (pip install cryptography)
"""
from __future__ import annotations

import argparse
import base64
import json
import os
import sys
import urllib.request

if hasattr(sys.stdout, "reconfigure"):
    sys.stdout.reconfigure(encoding="utf-8")

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
DEFAULT_PUB_KEY = os.path.join(ROOT, "deploy", "rsa-public-key-dev.pem")


def rsa_encrypt(plain: str, pub_key_path: str) -> str:
    try:
        from cryptography.hazmat.primitives import serialization
        from cryptography.hazmat.primitives.asymmetric import padding
    except ImportError:
        sys.exit("缺少依赖：pip install cryptography")

    with open(pub_key_path, "rb") as f:
        public_key = serialization.load_pem_public_key(f.read())

    # 必须与服务端 RsaCipher 的 "RSA/ECB/PKCS1Padding" 一致。
    # 换成 OAEP 会表现为「密码永远错」，两边都改才能换。
    cipher = public_key.encrypt(plain.encode("utf-8"), padding.PKCS1v15())
    return base64.b64encode(cipher).decode("ascii")


def post_json(url: str, payload: dict, timeout: int = 20) -> tuple[int, dict]:
    data = json.dumps(payload).encode("utf-8")
    req = urllib.request.Request(
        url, data=data,
        headers={"Content-Type": "application/json; charset=utf-8"},
        method="POST",
    )
    try:
        with urllib.request.urlopen(req, timeout=timeout) as resp:
            return resp.status, json.loads(resp.read().decode("utf-8"))
    except urllib.error.HTTPError as e:
        body = e.read().decode("utf-8")
        try:
            return e.code, json.loads(body)
        except json.JSONDecodeError:
            return e.code, {"raw": body}


def main() -> int:
    ap = argparse.ArgumentParser(description="RSA 加密登录，打印 token")
    ap.add_argument("username", nargs="?", default="admin")
    ap.add_argument("password", nargs="?", default="Admin@123")
    ap.add_argument("--base", default=os.environ.get("XT_BASE", "http://localhost:8080"))
    ap.add_argument("--tenant-code", default=None, help="账号在多个租户下重名时必填")
    ap.add_argument("--pub-key", default=DEFAULT_PUB_KEY)
    ap.add_argument("--json", action="store_true", help="打印完整响应而非仅 token")
    args = ap.parse_args()

    if not os.path.exists(args.pub_key):
        sys.exit(f"找不到公钥：{args.pub_key}")

    payload = {
        "username": args.username,
        "password": rsa_encrypt(args.password, args.pub_key),
    }
    if args.tenant_code:
        payload["tenantCode"] = args.tenant_code

    status, body = post_json(f"{args.base}/api/v1/auth/login", payload)

    if args.json:
        print(json.dumps(body, ensure_ascii=False, indent=2))
    elif body.get("code") == 0:
        d = body["data"]
        print(d["token"])
    else:
        print(f"[失败] HTTP {status}  code={body.get('code')}  msg={body.get('msg')}",
              file=sys.stderr)
        return 1

    return 0 if body.get("code") == 0 else 1


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