charging-cabinet/test/security-hardening/probe_export.py

#!/usr/bin/env python3
"""C 项补充实测：降权后（uid=pms）Excel 导出（写 ./exports）仍可用

流程：登录 → 创建导出任务 → 轮询任务状态 → 下载文件 → 打印文件落盘信息

退出码（仓库统一约定 0/1/2）：
- `0` = 导出链路可用（任务成功 + 下载到 xlsx）
- `1` = **判据失败**：导出链路真的坏了（任务 status=2 / 下载非 xlsx）
- `2` = **环境或夹具问题，本项"未验证"** —— 登录失败、或账号**缺少 `charge_record:export`
  权限码**（403）。⚠️ 这一档必须与 1 分开：本探针要回答的是"**降权后能不能往 exports 写文件**"，
  账号没权限根本走不到那一步，报 1 就是**把夹具缺陷说成产品缺陷**（误报与假绿同样有害）。
  实测（2026-09-20）测试库 `pms_dev` **只有 1 个账号且没有任何 user_roles**，
  因此这一项在补齐夹具前**恒为"未验证"**，不能算"导出不可用"。

夹具要求（两条路任选）：
1. 给测试账号配含 `charge_record:export` 的角色；或
2. 用环境变量给一个超管账号：`PMS_PROBE_PHONE` / `PMS_PROBE_PASSWORD`。
"""
import json
import os
import socket
import sys
import time
import urllib.error
import urllib.request
from urllib.parse import unquote, urlparse

BASE = "http://127.0.0.1:3000"
ENV_FILE = "/data/pms_dev/api-server/.env"
# 允许用环境变量换成"确有 charge_record:export 权限"的账号（默认沿用安全测试账号）
PHONE = os.environ.get("PMS_PROBE_PHONE", "13900000009")
PASSWORD = os.environ.get("PMS_PROBE_PASSWORD", "SecTest#2026")


class Redis:
    def __init__(self, url):
        u = urlparse(url)
        self.sock = socket.create_connection((u.hostname, u.port or 6379), timeout=5)
        self.fh = self.sock.makefile("rwb")
        if u.password:
            self.cmd("AUTH", unquote(u.password))
        self.cmd("SELECT", str(int((u.path or "/0").lstrip("/") or 0)))

    def _reply(self):
        line = self.fh.readline()
        if line.startswith(b"$"):
            n = int(line[1:])
            return None if n == -1 else self.fh.read(n + 2)[:-2].decode()
        return line.strip().decode()

    def cmd(self, *args):
        out = b"*%d\r\n" % len(args)
        for a in args:
            a = a.encode() if isinstance(a, str) else a
            out += b"$%d\r\n%s\r\n" % (len(a), a)
        self.fh.write(out)
        self.fh.flush()
        return self._reply()


def env(name):
    with open(ENV_FILE) as fh:
        for line in fh:
            if line.startswith(name + "="):
                return line.split("=", 1)[1].strip()
    raise SystemExit(f"{name} missing")


def http(method, path, data=None, headers=None):
    req = urllib.request.Request(BASE + path, data=data, method=method)
    for k, v in (headers or {}).items():
        req.add_header(k, v)
    try:
        with urllib.request.urlopen(req, timeout=30) as resp:
            return resp.status, resp.read()
    except urllib.error.HTTPError as exc:
        return exc.code, exc.read()


def main():
    redis = Redis(env("REDIS_URL"))
    st, body = http("GET", "/api/auth/challenge")
    cid = json.loads(body)["data"]["challenge_id"]
    x = redis.cmd("GET", f"challenge:{cid}")
    st, body = http(
        "POST",
        "/api/auth/login",
        json.dumps(
            {"phone": PHONE, "password": PASSWORD, "challenge_id": cid, "captcha_x": int(x)}
        ).encode(),
        {"Content-Type": "application/json"},
    )
    if st != 200:
        print(f"[login] {st} {body.decode()[:200]}")
        print("  ⚠️ 未验证（exit 2）：登录失败属环境/夹具问题，不是导出链路缺陷")
        return 2
    token = json.loads(body)["data"]["token"]
    auth = {"Authorization": f"Bearer {token}"}

    st, body = http("GET", "/api/charge-records/export?start_time=2020-01-01%2000:00:00&end_time=2030-01-01%2000:00:00", headers=auth)
    print(f"[export] create -> {st} {body.decode()[:200]}")
    if st == 403:
        # 关键区分：403 说明**账号没有权限码**，根本没走到"往 exports 写文件"那一步。
        # 把它算成失败（退出码 1）会让人以为"降权后导出不可用"，是误报。
        print("  ⚠️ 未验证（exit 2）：账号缺少 charge_record:export 权限，本项无法判定")
        print("     夹具要求见本文件顶部 docstring（配权限 或 设 PMS_PROBE_PHONE/PASSWORD）")
        return 2
    if st != 200:
        print(f"  ❌ exit 1：创建导出任务失败（HTTP {st}），非权限问题")
        return 1
    task_id = json.loads(body)["data"]["task_id"]

    # ⚠️ `download_tasks.status` 的**权威语义**（`db/migrate.rs:225` 列注释 +
    #    `routes/charge_records.rs:435` 的 `mark_task_completed` 写的就是 2）：
    #      0=排队  1=处理中  2=**完成**  3=失败
    # 本脚本此前**把语义搞反了**（等 `status == 1`、把 `status == 2` 判成"导出任务失败"），
    # 于是"导出成功"会被报成失败（假红），与 console 端 `DownloadCenter.tsx::statusText`
    # 和 `test/schema_alignment_e2e.py` 的口径都不一致。已按权威语义订正。
    STATUS_DONE, STATUS_FAILED = 2, 3

    file_path = None
    for _ in range(20):
        time.sleep(1)
        st, body = http("GET", f"/api/downloads/{task_id}", headers=auth)
        info = json.loads(body)["data"]
        cur = info.get("status")
        print(f"[export] task={task_id} status={cur} progress={info.get('progress')} file_path={info.get('file_path')}")
        if cur == STATUS_DONE:
            file_path = info.get("file_path")
            break
        if cur == STATUS_FAILED:
            print(f"  ❌ exit 1：导出任务失败（status=3）err={info.get('error_message')!r}")
            return 1
    else:
        print("  ❌ exit 1：120 秒内任务未进入终态（既没完成也没失败）—— 导出链路可能卡住")
        return 1

    if not file_path:
        print("  ❌ exit 1：任务标记完成但没有 file_path —— 落盘环节有问题")
        return 1

    st, blob = http("GET", f"/api/downloads/{task_id}/file", headers=auth)
    print(f"[download] GET /api/downloads/{task_id}/file -> {st} bytes={len(blob)} head={blob[:4]!r}")
    if st == 200 and blob[:2] == b"PK":
        print("  ✅ exit 0：导出链路可用（降权成 pms 后仍能往 exports/ 落盘并下载）")
        return 0
    print("  ❌ exit 1：下载失败或不是 xlsx（xlsx 是 zip，魔数应为 PK）")
    return 1


if __name__ == "__main__":
    sys.exit(main())