xiaotutongxue/scripts/e2e-m1-storage.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
对象存储 自检（七牛云 / 本地磁盘）。

    [x] 配置是否就绪（provider / bucket / domain / 密钥）
    [x] 材料上传 → 处理 → 落进对象存储
    [x] 取签名链接能下载到**字节一致**的内容        ← 真正的证明
    [x] **不带签名**的同地址被拒（证明空间确实是私有的，签名在起作用）
    [x] 对象键带目录前缀（与同空间其他项目隔离）

⚠️ 为什么要有这个脚本：
    "文件传上去了"和"文件**取得回来**"是两件事。七牛空间是私有的，
    少签一个名、签错有效期、键里的中文没编码对 —— 每一种都会让
    上传**成功**但下载失败（403/404）。只有真的把字节取回来比一遍才知道配对了。

跑法：
    python scripts/e2e-m1-storage.py
    没配置七牛时会直接跳过（exit 0），不会误报失败。
"""
from __future__ import annotations

import argparse
import base64
import hashlib
import hmac
import io
import json
import os
import re
import subprocess
import sys
import time
import urllib.error
import urllib.parse
import urllib.request
import uuid

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

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
SRC = os.path.join(ROOT, "backend", "xt-server", "src", "main", "resources")
LOCAL_YML = os.path.join(ROOT, "backend", "application-local.yml")
PG = os.environ.get("XT_PG_CONTAINER", "xt-pg")
UID = 992004
BASE = "http://localhost:8080"

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


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


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


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


def skip(n: str, why: str) -> None:
    print(f"  [SKIP] {n}\n         {why}")


def sql(q: str) -> str:
    r = subprocess.run(
        ["docker", "exec", "-e", "PGPASSWORD=fill_dev_pwd", PG, "psql", "-U", "fill",
         "-d", "filldb", "-t", "-A", "-c", q], capture_output=True, text=True, timeout=40)
    return (r.stdout or r.stderr).strip()


def read_yml(path: str) -> str:
    return open(path, encoding="utf-8").read() if os.path.exists(path) else ""


def resolve(v: str | None) -> str | None:
    """
    把 `${ENV:default}` 解析成实际值。

    ⚠️ 不解析的话自检会**说谎**：显示成 `provider=${XT_STORAGE_PROVIDER:local}` ——
    既不是 local 也不是 qiniu，读的人只会更糊涂。配置检查这一环一旦不可信，
    它给的"通过"也就没意义了。
    """
    if not v:
        return None
    m = re.fullmatch(r"\$\{([^:}]+):([^}]*)\}", v)
    if m:
        env = os.environ.get(m.group(1))
        return ((env if env is not None else m.group(2)).strip()) or None
    m = re.fullmatch(r"\$\{([^}]+)\}", v)
    if m:
        return (os.environ.get(m.group(1)) or "").strip() or None
    return v


def cfg_value(text: str, key: str) -> str | None:
    m = re.search(rf"^\s*{re.escape(key)}:\s*(.+?)\s*$", text, re.M)
    if not m:
        return None
    return resolve(m.group(1).strip().strip("'\"") or None)


def storage_cfg() -> dict:
    """
    读配置的顺序与 Spring 一致：**application-local.yml 覆盖 application-dev.yml**。
    不这么做会出现"我以为在用七牛、其实脚本读的是 dev 里的 local"这种自欺。
    """
    dev = read_yml(os.path.join(SRC, "application-dev.yml"))
    local = read_yml(LOCAL_YML)
    out = {}
    for k in ("provider", "prefix", "access-key", "secret-key", "bucket", "domain", "region"):
        out[k] = cfg_value(dev, k)
    for k in ("provider", "prefix", "access-key", "secret-key", "bucket", "domain", "region"):
        v = cfg_value(local, k)
        if v:
            out[k] = v
    return out


def b64u(b: bytes) -> str:
    return base64.urlsafe_b64encode(b).rstrip(b"=").decode()


def mint(secret: str, issuer: str, uid: int, tenant: int) -> str:
    key = base64.b64decode(secret)
    now = int(time.time())
    signing = (b64u(json.dumps({"alg": "HS512"}, separators=(",", ":")).encode()) + "."
               + b64u(json.dumps({"iss": issuer, "sub": str(uid), "tenantId": tenant,
                                  "orgIds": [], "roles": ["STUDENT"], "userType": 4,
                                  "iat": now, "exp": now + 3600}, separators=(",", ":")).encode()))
    return signing + "." + b64u(hmac.new(key, signing.encode(), hashlib.sha512).digest())


def call(method: str, path: str, token: str | None = None, body: dict | None = None) -> dict:
    h = {"Content-Type": "application/json; charset=utf-8"}
    if token:
        h["Authorization"] = f"Bearer {token}"
    data = json.dumps(body, ensure_ascii=False).encode() if body is not None else None
    req = urllib.request.Request(BASE + path, data=data, headers=h, method=method)
    try:
        with urllib.request.urlopen(req, timeout=90) as x:
            return json.loads(x.read().decode("utf-8"))
    except urllib.error.HTTPError as e:
        try:
            return json.loads(e.read().decode("utf-8"))
        except Exception:
            return {"_http": e.code}


def upload(token: str, sid: int, code: str, filename: str, data: bytes) -> dict:
    boundary = "----xte2e" + uuid.uuid4().hex[:12]
    parts = []
    for k, v in [("submissionId", str(sid)), ("materialCode", code)]:
        parts.append(f'--{boundary}\r\nContent-Disposition: form-data; name="{k}"\r\n\r\n{v}\r\n'.encode())
    parts.append(f'--{boundary}\r\nContent-Disposition: form-data; name="file"; filename="{filename}"\r\n'
                 f"Content-Type: image/jpeg\r\n\r\n".encode())
    parts.append(data)
    parts.append(f"\r\n--{boundary}--\r\n".encode())
    req = urllib.request.Request(
        BASE + "/api/v1/mp/files/upload", data=b"".join(parts), method="POST",
        headers={"Content-Type": f"multipart/form-data; boundary={boundary}",
                 "Authorization": f"Bearer {token}"})
    try:
        with urllib.request.urlopen(req, timeout=120) as x:
            return json.loads(x.read().decode("utf-8"))
    except urllib.error.HTTPError as e:
        try:
            return json.loads(e.read().decode("utf-8"))
        except Exception:
            return {"_http": e.code}


def http_get(url: str) -> tuple[int, bytes]:
    try:
        with urllib.request.urlopen(url, timeout=60) as x:
            return x.status, x.read()
    except urllib.error.HTTPError as e:
        return e.code, b""
    except Exception:
        return 0, b""


def small_jpeg() -> bytes:
    from PIL import Image
    buf = io.BytesIO()
    Image.new("RGB", (900, 1200), (206, 212, 218)).save(buf, "JPEG", quality=90)
    return buf.getvalue()


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--task", type=int, default=17)
    args = ap.parse_args()
    task_id = args.task

    print("=" * 72)
    print("对象存储 自检")
    print("=" * 72)

    cfg = storage_cfg()
    provider = (cfg.get("provider") or "local").lower()
    print(f"  provider = {provider}")
    print(f"  bucket   = {cfg.get('bucket')}")
    print(f"  domain   = {cfg.get('domain')}")
    print(f"  prefix   = {cfg.get('prefix')}")
    print(f"  region   = {cfg.get('region')}")

    if provider != "qiniu":
        print()
        skip("七牛云自检", "当前 provider=local（文件存在本地磁盘，多副本部署会丢文件）")
        print("""
  要启用七牛云：
    ① 编辑 backend/application-local.yml（已在 .gitignore 里，不会被提交）
    ② 取消 provider 那行的注释改成 qiniu，并填上 access-key / secret-key
    ③ 重启后端，再跑一次这个脚本
        """)
        print("=" * 72)
        print("跳过（不是失败）：当前在用本地磁盘")
        print("=" * 72)
        return 0

    for k in ("access-key", "secret-key", "bucket", "domain"):
        if not cfg.get(k):
            fail(f"配置 {k}", f"provider=qiniu 但 {k} 为空 —— 服务端会在启动时就拒绝")
            return 1
    ok("七牛配置齐全（provider/bucket/domain/密钥）")

    text = read_yml(os.path.join(SRC, "application-dev.yml")) + read_yml(os.path.join(SRC, "application.yml"))
    sec = re.search(r"^\s*secret:\s*(\S+)\s*$", text, re.M)
    iss = re.search(r"^\s*issuer:\s*(\S+)\s*$", text, re.M)
    if not sec or not iss:
        fail("读 JWT 配置", "找不到 secret/issuer")
        return 1
    tok = mint(sec.group(1), iss.group(1), UID, 1)

    rid = sql(f"SELECT id FROM collect_roster WHERE task_id={task_id} LIMIT 1")
    if not rid.isdigit():
        fail("取名册行", "任务下没有名册")
        return 1
    sql(f"UPDATE collect_roster SET bind_user_id={UID}, bind_type=1 WHERE id={rid}")

    r = call("POST", "/api/v1/mp/submissions", tok, {"taskId": task_id})
    if r.get("code") != 0:
        fail("建草稿", f"{r}")
        return 1
    sid = r["data"]["id"]
    print(f"  草稿 submission={sid}")

    try:
        print("\n[1] 上传 → 处理")
        img = small_jpeg()
        up = upload(tok, sid, "IDCARD", "身份证正面.jpg", img)
        if not check("上传成功（这一步已经把原图写进对象存储）", up.get("code") == 0, f"{up}"):
            return 1
        fid = up["data"]["id"]
        pr = call("POST", f"/api/v1/mp/files/{fid}/process", tok)
        check("压缩改名成功（规范件也写进去了）", pr.get("code") == 0, f"{pr}")

        print("\n[2] 签名链接")
        u = call("GET", f"/api/v1/mp/files/{fid}/url", tok)
        if not check("取到链接", u.get("code") == 0, f"{u}"):
            return 1
        d = u.get("data") or {}
        signed = d.get("url") or ""
        print(f"         {signed[:120]}{'…' if len(signed) > 120 else ''}")
        check("provider 是 qiniu（不是本地磁盘）", d.get("provider") == "qiniu", f"{d.get('provider')}")
        check("链接指向配置的域名", (cfg.get("domain") or "") in signed, f"domain={cfg.get('domain')}")
        check("是**签名**链接（带 e= 与 token=）",
              "e=" in signed and "token=" in signed, "没有签名参数，私有空间会 403")
        check("对象键带目录前缀", f"/{cfg.get('prefix')}/" in signed, f"prefix={cfg.get('prefix')}")
        check("有效期是正数秒", (d.get("expiresIn") or 0) > 0, f"{d.get('expiresIn')}")

        print("\n[3] 真的下载回来比一遍（这才是证明）")
        status, body = http_get(signed)
        check("签名链接返回 200", status == 200, f"HTTP {status}")
        if status == 200:
            check("下载到的内容非空", len(body) > 0, f"{len(body)} 字节")
            # 规范件是压缩过的，与上传的原图**不一定**等大（这是正常的）；
            # 这里比的是"取回来的确实是一张 JPEG 图"
            check("取回的是 JPEG（FFD8 开头）",
                  body[:2] == b"\xff\xd8", f"前两字节 {body[:2]!r}")

        print("\n[4] 不带签名应当被拒（证明空间确实是私有的）")
        unsigned = signed.split("?")[0]
        st2, _ = http_get(unsigned)
        check("裸地址被拒（401/403）", st2 in (401, 403), f"HTTP {st2} —— 如果返回 200 说明空间是公开的")

        print("\n[5] original=true 的行为（原始件默认不保留）")
        u2 = call("GET", f"/api/v1/mp/files/{fid}/url?original=true", tok)
        d2 = u2.get("data") or {}
        # 「原始件默认不保留」是客户的明确决定，所以这里**不该断言原始件存在** ——
        # 两种策略都是对的。要断言的是：**即使没有原始件，这个调用也不能报错**，
        # 而且要如实告诉我们取到的是哪一份。
        check("original=true 不报错", u2.get("code") == 0, f"{u2}")
        st3, body3 = http_get(d2.get("url") or "")
        check("取到的链接能下载", st3 == 200, f"HTTP {st3}")
        if d2.get("original") is True:
            check("开着 keep-raw：取回的就是原始件（字节一致）",
                  body3 == img, f"上传 {len(img)} / 取回 {len(body3)}")
        else:
            check("默认不留原始件：如实标注并退回到交付件",
                  len(body3) > 0 and body3[:2] == b"\xff\xd8", f"取了 {len(body3)} 字节")
    finally:
        print("\n[6] 清理")
        sql(f"DELETE FROM collect_file WHERE submission_id={sid}")
        sql(f"DELETE FROM collect_submission WHERE id={sid}")
        sql(f"UPDATE collect_roster SET bind_user_id=NULL, bind_time=NULL, bind_type=NULL WHERE id={rid}")
        ok("已清理草稿与名册绑定")
        print("""
  ⚠️ 对象存储里的文件**没有删** —— 那正是要的：原始件留着备查，
     认定机构退件时要能回答"是不是你们压坏的"。要清就从七牛控制台按前缀清。""")

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


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