产品介绍

#!/usr/bin/env python3
"""
JXWAF 安全模型服务 Python 客户端
支持单条/批量检测、自动轮询等待分析结果
"""
import os
import sys
import json
import time
import argparse
from pathlib import Path

import requests


class JxwafClient:
    def __init__(self, api_key=None, base_url=None):
        self.api_key = api_key or os.environ.get("JXWAF_API_KEY", "")
        self.base_url = (base_url or os.environ.get("JXWAF_BASE_URL", "")).rstrip("/")
        if not self.api_key:
            raise ValueError("API Key 未设置，请通过参数或环境变量 JXWAF_API_KEY 配置")
        if not self.base_url:
            raise ValueError("服务地址未设置，请通过参数或环境变量 JXWAF_BASE_URL 配置")
        self.session = requests.Session()
        self.session.headers.update({
            "Content-Type": "application/json",
            "jxwaf-model-auth": self.api_key,
        })

    def _submit(self, payload: str) -> dict:
        resp = self.session.post(
            f"{self.base_url}/jxwaf_model/web_attack_check",
            data=payload,
            timeout=120,
        )
        resp.raise_for_status()
        return resp.json()

    def check(self, payload: str, poll_interval: int = 2, max_wait: int = 300) -> dict:
        """提交检测并等待分析完成。"""
        submitted = False
        deadline = time.time() + max_wait

        while time.time() < deadline:
            result = self._submit(payload)
            if result.get("status") == "error":
                return result

            verdict = result.get("result", "")
            if verdict in ("safe", "attack"):
                return result

            # 首次提交后，后续使用 token 查询加速命中缓存
            if not submitted:
                submitted = True
                payload = json.dumps({"token": result.get("token", "")})

            time.sleep(poll_interval)

        return result

    def check_batch(self, items: list, poll_interval: int = 2, max_wait: int = 300, verbose: bool = False) -> list:
        """批量检测，支持实时进度输出。items 可以是字符串列表或 (标签, 数据) 元组列表。"""
        results = []
        for i, item in enumerate(items):
            if isinstance(item, (list, tuple)):
                label, payload = item
            else:
                label, payload = None, item

            if verbose:
                tag = f"[{label}] " if label else ""
                print(f"[{i+1}/{len(items)}] {tag}提交检测...", end=" ", flush=True)

            r = self.check(str(payload), poll_interval, max_wait)
            results.append(r)

            if verbose:
                verdict = r.get("result", "?")
                attack_type = r.get("attack_type", [])
                if verdict == "attack":
                    print(f"⚠ 攻击 ({', '.join(attack_type)})")
                elif verdict == "safe":
                    print("✓ 安全")
                elif verdict == "unknown":
                    print("⏳ 超时未完成")
                else:
                    print(f"✗ {verdict}")

        return results


def load_payloads(source: str) -> list[str]:
    """从文件加载待检测数据，每行一条纯文本请求。以 '#' 开头的行视为注释，空行自动忽略。"""
    path = Path(source)
    if not path.exists():
        raise FileNotFoundError(f"文件不存在: {source}")
    payloads = []
    with open(path, "r", encoding="utf-8") as f:
        for line in f:
            line = line.strip()
            if not line or line.startswith("#"):
                continue
            payloads.append(line)
    return payloads


def format_result(result: dict) -> str:
    """格式化检测结果为可读字符串。"""
    verdict = result.get("result", "?")
    token = result.get("token", "")
    model = result.get("ai_model", "")
    attack_type = result.get("attack_type", [])
    if verdict == "attack":
        return f"⚠ 攻击 | 类型: {', '.join(attack_type)} | 模型: {model} | Token: {token}"
    elif verdict == "safe":
        return f"✓ 安全 | 模型: {model} | Token: {token}"
    elif verdict == "unknown":
        return f"⏳ 待分析 | Token: {token}"
    else:
        return f"✗ {verdict} | Token: {token}"


def main():
    parser = argparse.ArgumentParser(description="JXWAF安全模型服务 - Web安全检测客户端")
    parser.add_argument("payload", nargs="?", help="待检测数据（也可从 stdin 读取）")
    parser.add_argument("-k", "--api-key", default=os.environ.get("JXWAF_API_KEY", ""),
                        help="API Key（可通过环境变量 JXWAF_API_KEY 设置）")
    parser.add_argument("-s", "--server", default=os.environ.get("JXWAF_BASE_URL", ""),
                        help="服务地址（可通过环境变量 JXWAF_BASE_URL 设置）")
    parser.add_argument("-f", "--file", help="从文件读取批量数据（每行一条纯文本请求）")
    parser.add_argument("-o", "--output", help="结果输出文件（JSON 格式）")
    parser.add_argument("-i", "--interval", type=int, default=2, help="轮询间隔秒数（默认 2）")
    parser.add_argument("-w", "--max-wait", type=int, default=300, help="单条最大等待秒数（默认 300）")
    parser.add_argument("-v", "--verbose", action="store_true", help="打印详细过程")
    args = parser.parse_args()

    client = JxwafClient(api_key=args.api_key, base_url=args.server)

    # 批量检测模式
    if args.file:
        items = load_payloads(args.file)
        if not items:
            print("文件中没有待检测数据", file=sys.stderr)
            sys.exit(1)
        if args.verbose:
            print(f"从文件加载 {len(items)} 条数据\n")
        results = client.check_batch(items, args.interval, args.max_wait, verbose=args.verbose)

        print(f"\n===== 检测结果 ({len(results)} 条) =====")
        attack_count = sum(1 for r in results if r.get("result") == "attack")
        safe_count = sum(1 for r in results if r.get("result") == "safe")
        unknown_count = sum(1 for r in results if r.get("result") == "unknown")
        print(f"攻击: {attack_count}  安全: {safe_count}  待分析: {unknown_count}\n")

        for i, (item, r) in enumerate(zip(items, results)):
            # 纯文本模式下直接显示序号
            print(f"[{i+1}] " + format_result(r))

        if args.output:
            with open(args.output, "w", encoding="utf-8") as f:
                json.dump(results, f, ensure_ascii=False, indent=2)
            print(f"\n结果已保存至: {args.output}")
        return

    # 单条检测模式
    payload = args.payload
    if not payload:
        if not sys.stdin.isatty():
            payload = sys.stdin.read().strip()
        else:
            parser.print_help()
            sys.exit(1)

    r = client.check(payload, args.interval, args.max_wait)
    print(format_result(r))
    if args.output:
        with open(args.output, "w", encoding="utf-8") as f:
            json.dump(r, f, ensure_ascii=False, indent=2)


if __name__ == "__main__":
    main()