#!/usr/bin/env python3
"""B 项实测探针（直连测试后端 127.0.0.1:3000，X-Real-IP 用于隔离不同"客户端"桶）

1. 成功登录清零：2 次错密码 → 计数=2 → 正确密码登录成功 → 计数键被删除
2. 锁定：5 次错密码 → 第 6 次（以及锁定后的正确密码）都返回统一文案 401
3. challenge 限流：同一 IP 连打 12 次 → 第 11/12 次 429
4. 并发闸门：N 个不同 IP 并发请求 challenge → 观察是否被排队/超时（Semaphore=2）
"""
import concurrent.futures
import json
import socket
import sys
import time
import urllib.error
import urllib.request
from urllib.parse import unquote, urlparse

from verdict import Verdict   # 同目录：判据 + 统一退出码（0 通过 / 1 有失败 / 2 环境问题）

BASE = "http://127.0.0.1:3000"
ENV_FILE = "/data/pms_dev/api-server/.env"
PHONE = "13900000009"
GOOD = "SecTest#2026"
BAD = "WrongPass#1"
FAIL_KEY = f"auth:fail:phone:{PHONE}"


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


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

    def _reply(self):
        line = self.fh.readline()
        if line.startswith(b"+"):
            return line[1:].strip().decode()
        if line.startswith(b"-"):
            raise RuntimeError(line.decode().strip())
        if line.startswith(b":"):
            return int(line[1:])
        if line.startswith(b"$"):
            n = int(line[1:])
            if n == -1:
                return None
            return self.fh.read(n + 2)[:-2].decode()
        raise RuntimeError(f"unexpected reply: {line!r}")

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


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=20) as resp:
            return resp.status, resp.read()
    except urllib.error.HTTPError as exc:
        return exc.code, exc.read()


def login_attempt(redis, ip, password):
    """一次完整登录：取 challenge → （从 Redis 读正确 x，模拟人正确拖动）→ 提交"""
    st, body = http("GET", "/api/auth/challenge", headers={"X-Real-IP": ip})
    if st != 200:
        return st, body
    cid = json.loads(body)["data"]["challenge_id"]
    correct_x = redis.cmd("GET", f"challenge:{cid}")
    payload = json.dumps(
        {"phone": PHONE, "password": password, "challenge_id": cid, "captcha_x": int(correct_x)}
    ).encode()
    return http("POST", "/api/auth/login", payload, {"Content-Type": "application/json", "X-Real-IP": ip})


def counters(redis, ip):
    return {
        "phone": redis.cmd("GET", FAIL_KEY),
        "ip": redis.cmd("GET", f"auth:fail:ip:{ip}"),
        "phone_ttl": redis.cmd("TTL", FAIL_KEY),
    }


def main():
    # 判据容器：失败 → exit 1（此前写死 return 0，属假绿）
    v = Verdict("probe_bruteforce（登录锁定 / challenge 限流）")
    redis = Redis(read_env("REDIS_URL"))
    for k in (FAIL_KEY, "auth:fail:ip:10.9.9.4", "auth:fail:ip:10.9.9.5", "auth:fail:ip:10.9.9.6"):
        redis.cmd("DEL", k)

    print("===== B1 成功登录清零失败计数 =====")
    for i in (1, 2):
        st, body = login_attempt(redis, "10.9.9.4", BAD)
        print(f"  错密码#{i} status={st} body={body.decode()[:120]}")
    c_after_fail = counters(redis, "10.9.9.4")
    print(f"  计数（2 次失败后）: {c_after_fail}")
    v.check("B1a 2 次错密码应记入失败计数", c_after_fail["phone"] is not None,
            f'phone={c_after_fail["phone"]}')
    st, body = login_attempt(redis, "10.9.9.4", GOOD)
    print(f"  正确密码 status={st} body={body.decode()[:120]}")
    v.check("B1b 正确密码应登录成功", st == 200, f"status={st}")
    c_after_ok = counters(redis, "10.9.9.4")
    print(f"  计数（登录成功后应为 None）: {c_after_ok}")
    v.check("B1c 登录成功后失败计数应被清零", c_after_ok["phone"] is None,
            f'phone={c_after_ok["phone"]}')

    print("\n===== B2 连续失败锁定（阈值 5 次 / 15 分钟）=====")
    last_st = None
    for i in range(1, 7):
        st, body = login_attempt(redis, "10.9.9.5", BAD)
        last_st = st
        print(f"  错密码#{i} status={st} body={body.decode()[:120]}  counters={counters(redis, '10.9.9.5')}")
    v.check("B2a 第 6 次错密码应被拒（非 200）", last_st != 200, f"status={last_st}")
    print("  锁定后再用【正确密码】尝试（应仍被拒，且文案与密码错误完全一致，不泄露锁定状态）:")
    st, body = login_attempt(redis, "10.9.9.7", GOOD)
    print(f"  正确密码 status={st} body={body.decode()[:120]}")
    v.check("B2b 锁定后用正确密码也应被拒（文案不泄露锁定状态）", st != 200, f"status={st}")

    print("\n===== B3 challenge 限流（同一 IP 12 次）=====")
    for i in range(1, 13):
        st, body = http("GET", "/api/auth/challenge", headers={"X-Real-IP": "10.9.9.6"})
        print(f"  challenge#{i} status={st} body={body.decode()[:80]}")
    # ⚠️ 不加硬判据：触发点取决于 nginx `auth_challenge_zone` 的 rate（AGENTS.md 记 20r/m，
    #    而本脚本注释写"第 11 次起 429"）—— 口径不一致时写死判据会误报，故保持观察输出。

    print("\n===== B4 并发闸门（Semaphore=2，不同 IP 并发请求 challenge）=====")
    t0 = time.time()
    st, _ = http("GET", "/api/auth/challenge", headers={"X-Real-IP": "10.9.7.1"})
    single = time.time() - t0
    print(f"  单次 challenge 耗时 {single*1000:.0f} ms")
    n = 40
    t0 = time.time()
    with concurrent.futures.ThreadPoolExecutor(max_workers=n) as pool:
        results = list(
            pool.map(
                lambda i: http("GET", "/api/auth/challenge", headers={"X-Real-IP": f"10.9.8.{i}"})[0],
                range(1, n + 1),
            )
        )
    wall = time.time() - t0
    codes = {}
    for r in results:
        codes[r] = codes.get(r, 0) + 1
    print(f"  {n} 并发（不同 IP）耗时 {wall:.2f}s，状态码分布={codes}")
    # ⚠️ 同上：并发耗时/状态码只作观察，不设判据。

    print("\n===== 清理测试计数键 =====")
    for k in (FAIL_KEY, "auth:fail:ip:10.9.9.4", "auth:fail:ip:10.9.9.5", "auth:fail:ip:10.9.9.7"):
        redis.cmd("DEL", k)
    print(f"  清理后 phone 计数 = {redis.cmd('GET', FAIL_KEY)}")
    return v.exit_code()      # 0=判据全过，1=有判据失败


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