xiaotutongxue/scripts/e2e-m1-tenant-domain.py

# -*- coding: utf-8 -*-
"""
租户域名 E2E：**用子域前缀判断"这是哪个单位的站点"**（t1.xxx.com / t2.xxx.com）。

验的核心只有两条，但它们决定了这个功能能不能上：

  ① **登录时不用再填租户码**。在此之前，「这次是哪个单位」要么靠人手填租户码
     （账号跨单位重名时报 40012），要么靠扫码 scene —— 而域名本身就是最自然的答案：
     用户在 t1.xxx.com 上打开页面时，系统不该再问他"你是哪个单位的"。

  ② ⚠️ **Host 只用于"定位"，授权一律以 token 为准，且两者不一致必须拒绝**。
     不校验的话：在 t1.xxx.com 上拿着 t2 的 token 就能读 t2 的数据 ——
     界面一切正常（前端本来就会带 token），是个安静且极难被发现的越权。
     所以这条比第 ① 条更重要。

跑法：python scripts/e2e-m1-tenant-domain.py
前提：后端 8080 在跑，且 `fill.tenant-domain` 已启用（dev 默认
      enabled=true / base=localhost —— 只有 *.localhost 会走这条路，不影响其它域名）。
"""
from __future__ import annotations

import argparse
import base64
import http.client
import json
import os
import sys

sys.stdout.reconfigure(encoding="utf-8")

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

PASSED: list[str] = []
FAILED: list[tuple[str, str]] = []


def ok(name: str) -> None:
    PASSED.append(name)
    print(f"  [OK]   {name}")


def fail(name: str, detail: str) -> None:
    FAILED.append((name, detail))
    print(f"  [FAIL] {name}\n         {detail}")


def check(name: str, cond: bool, detail: str = "") -> bool:
    if cond:
        ok(name)
    else:
        fail(name, detail or "断言不成立")
    return cond


def rsa_encrypt(plain: str) -> str:
    from cryptography.hazmat.primitives import serialization
    from cryptography.hazmat.primitives.asymmetric import padding
    with open(PUB_KEY, "rb") as f:
        pub = serialization.load_pem_public_key(f.read())
    return base64.b64encode(pub.encrypt(plain.encode(), padding.PKCS1v15())).decode()


def call(method: str, path: str, host: str, token: str | None = None,
         body: dict | None = None) -> dict:
    """
    发一个请求，**Host 由调用方指定**。

    用 http.client 而不是 urllib：urllib 会按 URL 自己算 Host，
    而这里恰恰要伪造 Host（模拟 t1.localhost 这种子域），必须完全掌控这个头。
    """
    payload = json.dumps(body, ensure_ascii=False).encode("utf-8") if body is not None else None
    headers = {"Host": host, "Content-Type": "application/json; charset=utf-8"}
    if token:
        headers["Authorization"] = f"Bearer {token}"
    conn = http.client.HTTPConnection("127.0.0.1", 8080, timeout=30)
    try:
        conn.request(method, path, body=payload, headers=headers)
        resp = conn.getresponse()
        raw = resp.read().decode("utf-8")
        try:
            return json.loads(raw)
        except json.JSONDecodeError:
            return {"_raw": raw[:200], "_http": resp.status}
    finally:
        conn.close()


def login(host: str, tenant_code: str | None = None, account: str = "admin",
          password: str = "Admin@123") -> dict:
    body: dict = {"username": account, "password": rsa_encrypt(password)}
    if tenant_code is not None:
        body["tenantCode"] = tenant_code
    return call("POST", "/api/v1/auth/login", host, body=body)


def main() -> int:
    ap = argparse.ArgumentParser(description="租户域名 E2E")
    ap.add_argument("--password", default="Admin@123")
    args = ap.parse_args()

    print("=" * 76)
    print("租户域名 E2E（t{tenantId}.{base}）")
    print("=" * 76)

    # ── 1. 域名就是租户身份：不用填租户码 ─────────────────
    print("\n[1] 在 t1.localhost 上，不填租户码也应能登录")
    # 先证明"没有域名时确实需要租户码"——否则下面的成功说明不了是域名起的作用
    no_domain = login("localhost", None)
    check("对照：不带域名 + 不填租户码 → 40012（账号跨单位重名）",
          no_domain.get("code") == 40012, f"{no_domain}")

    with_domain = login("t1.localhost", None)
    check("**t1.localhost + 不填租户码 → 登录成功**（域名提供了租户身份）",
          with_domain.get("code") == 0, f"{with_domain}")
    tenant = (with_domain.get("data") or {}).get("tenant") or {}
    check("登进去的确实是租户 1（demo）", tenant.get("id") == 1 and tenant.get("tenantCode") == "demo",
          f"{tenant}")
    token1 = (with_domain.get("data") or {}).get("token")

    # 带端口、带大写、多一层子域，都应解析到同一个租户
    check("Host 带端口（t1.localhost:8080）也能识别",
          login("t1.localhost:8080", None).get("code") == 0, "带端口失败")
    check("大小写不敏感（T1.LocalHost）",
          login("T1.LocalHost", None).get("code") == 0, "大小写失败")
    check("多一层子域（a.t1.localhost）取最后一段",
          login("a.t1.localhost", None).get("code") == 0, "多层子域失败")

    # ── 2. 与手填租户码矛盾时：明确报错，不静默忽略 ────────
    print("\n[2] 域名与手填的租户码矛盾时必须报错")
    conflict = login("t1.localhost", "somewrongcode")
    check("t1 域名 + 别的租户码 → 明确报错（不静默忽略）",
          conflict.get("code") == 40001, f"{conflict}")
    check("错误里说清了「当前访问的是哪个单位」与「不一致」",
          "不一致" in (conflict.get("msg") or "") and "站点" in (conflict.get("msg") or ""),
          f"{conflict.get('msg')}")

    # 一致的租户码是允许的（老客户端仍会带上）
    same = login("t1.localhost", "demo")
    check("t1 域名 + 正确的租户码 → 正常登录", same.get("code") == 0, f"{same}")

    # ── 3. ★ Host 只用于定位，授权以 token 为准 ────────────
    print("\n[3] ★ Host 与 token 的租户不一致 → 必须拒绝")
    check("拿 t1 的 token 访问 t1 域名 → 正常", 
          call("GET", "/api/v1/users?pageNum=1&pageSize=1", "t1.localhost", token1).get("code") == 0,
          "同租户被误拒")
    # 14 号租户是 e2e 建的（status=1，存在且启用），所以 t14.localhost 能被解析成一个真实租户
    cross = call("GET", "/api/v1/users?pageNum=1&pageSize=1", "t14.localhost", token1)
    check("**拿 t1 的 token 访问 t14 域名 → 403（拒绝）**",
          cross.get("code") == 40301 or cross.get("_http") == 403
          or "另一个单位" in (cross.get("msg") or ""), f"{cross}")
    check("拒绝时给出可理解的提示（不是一句「无权限」）",
          "另一个单位" in (cross.get("msg") or ""), f"{cross.get('msg')}")

    # ── 4. 边界：不合法/不存在的前缀一律回退到老路径 ────────
    print("\n[4] 边界：解析不出来就当「没有租户域名」，回退到老路径")
    check("前缀不是 t{数字}（x1.localhost）→ 回退（仍 40012）",
          login("x1.localhost", None).get("code") == 40012, "非 t 前缀被误当租户")
    check("租户不存在（t999999.localhost）→ 回退（仍 40012，而不是「找不到此人」）",
          login("t999999.localhost", None).get("code") == 40012,
          f"{login('t999999.localhost', None)}")
    check("主域名本身（localhost）= 平台入口，不属于任何租户",
          login("localhost", "demo").get("code") == 0, "主域名 + 租户码被误拒")
    check("无关域名（example.com）不解析",
          login("example.com", None).get("code") == 40012, "无关域名被误当租户")

    print()
    print("=" * 76)
    print(f"通过 {len(PASSED)}　失败 {len(FAILED)}　共 {len(PASSED) + len(FAILED)}")
    if FAILED:
        for n, det in FAILED:
            print(f"  · {n}\n      {det}")
    print("=" * 76)
    return 1 if FAILED else 0


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