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

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

流程：登录 → 创建导出任务 → 轮询任务状态 → 下载文件 → 打印文件落盘信息
"""
import json
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"
PHONE = "13900000009"
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]}")
        return 1
    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 != 200:
        return 1
    task_id = json.loads(body)["data"]["task_id"]

    file_path = None
    for _ in range(15):
        time.sleep(1)
        st, body = http("GET", f"/api/downloads/{task_id}", headers=auth)
        info = json.loads(body)["data"]
        print(f"[export] task={task_id} status={info.get('status')} progress={info.get('progress')} file_path={info.get('file_path')}")
        if info.get("status") == 1:
            file_path = info.get("file_path")
            break
        if info.get("status") == 2:
            print("导出任务失败")
            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}")
    return 0 if (st == 200 and blob[:2] == b"PK") else 1


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