#!/usr/bin/env python3
"""
Session / Token Boundary Lab Tool

LAB ONLY | AUTHORIZED TARGETS ONLY
"""

from __future__ import annotations

import argparse
import re
import sys
from pathlib import Path
from typing import Any, Dict, List

import requests

SCRIPTS_DIR = Path(__file__).resolve().parents[3] / "scripts"
if str(SCRIPTS_DIR) not in sys.path:
    sys.path.insert(0, str(SCRIPTS_DIR))

from tool_contract import add_common_args, emit_report, ensure_authorized, make_report, parse_headers, write_evidence  # noqa: E402


COOKIE_ATTRS = ["HttpOnly", "Secure", "SameSite", "Path", "Domain"]
STORAGE_PATTERNS = {
    "localStorage": re.compile(r"localStorage\.(setItem|getItem)|window\.localStorage", re.I),
    "sessionStorage": re.compile(r"sessionStorage\.(setItem|getItem)|window\.sessionStorage", re.I),
    "token-ish": re.compile(r"(jwt|token|authorization|bearer)", re.I),
}


def analyze(target: str, timeout: float, headers: Dict[str, str]) -> Dict[str, Any]:
    response = requests.get(target, timeout=timeout, headers=headers, verify=False)
    cookies = []
    for raw in response.headers.get("Set-Cookie", "").split(","):
        raw = raw.strip()
        if not raw:
            continue
        attrs = {attr: (attr.lower() in raw.lower()) for attr in COOKIE_ATTRS}
        cookies.append({"raw": raw[:300], "attributes": attrs})
    storage_hits = []
    for name, pattern in STORAGE_PATTERNS.items():
        if pattern.search(response.text):
            storage_hits.append(name)
    suspicious_headers = []
    for name in ["Set-Cookie", "Authorization", "X-Forwarded-User", "X-Original-URL"]:
        if response.headers.get(name):
            suspicious_headers.append({"name": name, "value": response.headers.get(name)[:200]})
    return {
        "status_code": response.status_code,
        "cookies": cookies,
        "storage_hits": storage_hits,
        "suspicious_headers": suspicious_headers,
        "body_excerpt": response.text[:600],
    }


def main() -> int:
    parser = argparse.ArgumentParser(description="Session / Token Boundary Lab Tool")
    parser.add_argument("--target", required=True, help="目标 URL")
    parser.add_argument("--timeout", type=float, default=8.0, help="请求超时时间")
    add_common_args(parser)
    args = parser.parse_args()
    ensure_authorized(args, parser)

    headers = parse_headers(args.header)
    findings = analyze(args.target, args.timeout, headers)
    evidence_refs = []
    ref = write_evidence(args, "session-lab.json", findings)
    if ref:
        evidence_refs.append(ref)
    suspicious = len(findings["cookies"]) + len(findings["storage_hits"]) + len(findings["suspicious_headers"])
    report = make_report(
        tool="session-lab",
        mode="cookie-storage-session-boundary-check",
        target=args.target,
        status="verified" if suspicious else "needs-review",
        severity="medium" if suspicious else "info",
        payload_or_probe=findings,
        request_summary={"timeout": args.timeout, "header_names": sorted(headers.keys())},
        evidence_refs=evidence_refs,
        destructive_risk="low",
        args=args,
    )
    text_lines = [
        "=" * 60,
        "Session / Token Boundary Lab Tool",
        "=" * 60,
        f"Target: {args.target}",
        f"Cookie Findings: {len(findings['cookies'])}",
        f"Storage Hits: {len(findings['storage_hits'])}",
        f"Suspicious Headers: {len(findings['suspicious_headers'])}",
    ]
    return emit_report(args, report, text_lines)


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