Skip to content

L
LLM-Filter
Commit 6e2b6222 authored by uuo's avatar uuo
Browse files
Options

refactor(security-service): 重构Zabbix服务异步调用并简化API端点 

- 在ZabbixService中添加_run_blocking方法处理阻塞调用兼容性
- 将collect_device_data等方法的同步调用改为异步，避免事件循环阻塞
- 简化API端点，将数据库日志记录逻辑移至SecurityService内部
- 统一时间处理，修复datetime.timezone.utc引用错误
- 重构测试脚本为多功能工具，支持Zabbix数据采集和API调用
- 优化应用启动逻辑，改进健康检查端点实现
parent 4d105058
Hide whitespace changes
Inline Side-by-side
Showing with 343 additions and 239 deletions
microservices/security-service/app/api/v1/endpoints.py
View file @ 6e2b6222
from datetime import datetime
from fastapi import APIRouter, Depends, Query
from app.schemas.payloads import *
from app.services.analysis import SecurityService
from app.services.rss import RSSService
from app.services.zabbix_service import ZabbixService
from app.core.security import get_current_admin
from app.core.database import db
from datetime import datetime, timezone
router = APIRouter()
zabbix_service = ZabbixService()
......@@ -14,15 +14,7 @@ rss_service = RSSService()
@router.post("/analysis", response_model=SecurityAnalysisResponse)
async def analyze_risks(request: SecurityAnalysisRequest, admin: dict = Depends(get_current_admin)):
result = await service.analyze_risks(request.devices)
# 异步存储结果到 MongoDB
if db.db is not None:
log_entry = result.model_dump()
log_entry["created_at"] = datetime.now(timezone.utc)
await db.db.security_analysis_logs.insert_one(log_entry)
return result
return await service.analyze_risks(request.devices)
@router.get("/analysis/history", response_model=HistoryQueryResponse)
async def get_analysis_history(
......@@ -38,15 +30,12 @@ async def get_analysis_history(
@router.post("/attack-advice", response_model=AttackAdviceResponse)
async def get_attack_advice(request: AttackAdviceRequest, admin: dict = Depends(get_current_admin)):
result = await service.get_attack_advice(request.attack_type, request.target_device, request.logs)
# 异步存储结果到 MongoDB
if db.db is not None:
log_entry = result.model_dump()
log_entry["created_at"] = datetime.now(timezone.utc)
await db.db.attack_advice_logs.insert_one(log_entry)
return result
return await service.get_attack_advice(
attack_type=request.attack_type,
target=request.target_device,
logs=request.logs,
severity=request.severity,
)
@router.get("/attack-advice/history", response_model=HistoryQueryResponse)
async def get_attack_advice_history(
......
microservices/security-service/app/main.py
View file @ 6e2b6222
from fastapi import FastAPI, HTTPException
from contextlib import asynccontextmanager
from datetime import datetime
import logging
from fastapi import FastAPI, HTTPException
from app.api.v1.endpoints import router as security_router
from app.core.config import settings
from app.core.database import db
from datetime import datetime
import logging
import logging
from app.services.zabbix_service import ZabbixService
logger = logging.getLogger(__name__)
@asynccontextmanager
async def lifespan(app: FastAPI):
# Startup
logger = logging.getLogger(__name__)
db.connect()
try:
zabbix_service = ZabbixService()
sync_status = zabbix_service.get_sync_status()
if sync_status.get("collector_initialized"):
logger.info("Zabbix服务初始化成功")
else:
logger.warning("Zabbix服务初始化失败，请检查Zabbix配置")
except Exception as e:
logger.error(f"Zabbix服务连接检查失败: {e}")
yield
db.close()
app = FastAPI(title=settings.PROJECT_NAME, version="2.0.0", lifespan=lifespan)
@app.get("/")
def health_check():
"""根路径健康检查"""
return {"status": "ok", "service": "security-service", "version": "2.0.0"}
@app.get("/health")
async def detailed_health_check():
"""详细健康检查端点"""
try:
# 检查数据库连接
db_status = "ok" if db.db else "error"
# 检查Zabbix服务
zabbix_status = {"status": "not_configured", "message": "Zabbix服务未配置"}
zabbix_status = {"status": "not_configured"}
try:
from app.services.zabbix_service import ZabbixService
zabbix_service = ZabbixService()
sync_status = zabbix_service.get_sync_status()
if sync_status["collector_initialized"]:
if sync_status.get("collector_initialized"):
zabbix_status = {
"status": "ok",
"last_sync": sync_status["last_sync_time"]
"last_sync": sync_status.get("last_sync_time"),
}
else:
zabbix_status = {
"status": "error",
"message": "Zabbix collector未初始化"
"message": "Zabbix collector未初始化",
}
except Exception as e:
zabbix_status = {
"status": "error",
"message": str(e)
"message": str(e),
}
from datetime import datetime
return {
"status": "healthy",
"timestamp": datetime.now().isoformat(),
......@@ -54,63 +71,34 @@ async def detailed_health_check():
"version": "2.0.0",
"components": {
"database": {"status": db_status},
"zabbix": zabbix_status
}
"zabbix": zabbix_status,
},
}
except Exception as e:
logger.error(f"Health check failed: {e}")
raise HTTPException(status_code=503, detail=f"Service unhealthy: {str(e)}")
@app.get("/ready")
async def readiness_check():
"""就绪检查端点"""
try:
# 检查关键依赖是否就绪
if not db.db:
return {
"status": "not_ready",
"timestamp": datetime.now().isoformat(),
"reason": "database_not_connected"
"reason": "database_not_connected",
}
return {
"status": "ready",
"timestamp": datetime.now().isoformat()
"timestamp": datetime.now().isoformat(),
}
except Exception as e:
return {
"status": "not_ready",
"timestamp": datetime.now().isoformat(),
"reason": str(e)
"reason": str(e),
}
@asynccontextmanager
async def lifespan(app: FastAPI):
# Startup
logger = logging.getLogger(__name__)
# 连接数据库
db.connect()
# 初始化Zabbix服务连接检查
try:
from app.services.zabbix_service import ZabbixService
zabbix_service = ZabbixService()
sync_status = zabbix_service.get_sync_status()
if sync_status["collector_initialized"]:
logger.info("✅ Zabbix服务初始化成功")
else:
logger.warning("⚠️ Zabbix服务初始化失败，请检查Zabbix配置")
except Exception as e:
logger.error(f"❌ Zabbix服务连接检查失败: {e}")
yield
# Shutdown
db.close()
app = FastAPI(title=settings.PROJECT_NAME, lifespan=lifespan)
# 注册路由
app.include_router(security_router, prefix=f"{settings.API_V1_STR}/security", tags=["Security"])
microservices/security-service/app/services/analysis.py
View file @ 6e2b6222
import json
import httpx
from datetime import datetime
from datetime import datetime, timezone
from typing import List, Dict, Any, Optional
from pydantic import ValidationError
from app.schemas.payloads import *
......@@ -67,7 +67,7 @@ class SecurityService:
try:
if db.db is not None:
log_entry = result.model_dump()
log_entry["created_at"] = datetime.now(datetime.timezone.utc)
log_entry["created_at"] = datetime.now(timezone.utc)
log_entry["device_count"] = len(devices)
await db.db.security_analysis_logs.insert_one(log_entry)
logger.info("安全分析结果已保存到MongoDB")
......@@ -76,23 +76,41 @@ class SecurityService:
return result
async def get_attack_advice(self, attack_type: str, target: str, logs: str) -> AttackAdviceResponse:
async def get_attack_advice(self, attack_type: str, target: str, logs: str, severity: Optional[str] = None) -> AttackAdviceResponse:
"""
任务：攻击应急建议
"""
# 构造结构化数据
data = {
"attack_type": attack_type,
"target_device": target,
"logs": logs
"severity": severity,
"logs": logs,
}
inputs = {
"task_type": "advice",
"context_data": json.dumps(data, ensure_ascii=False)
"context_data": json.dumps(data, ensure_ascii=False),
}
return await self._call_llm(inputs, AttackAdviceResponse)
result = await self._call_llm(inputs, AttackAdviceResponse)
try:
if db.db is not None:
log_entry = result.model_dump()
log_entry.update(
{
"created_at": datetime.now(timezone.utc),
"attack_type": attack_type,
"target_device": target,
"severity": severity,
}
)
await db.db.attack_advice_logs.insert_one(log_entry)
logger.info("攻击建议结果已保存到MongoDB")
except Exception as e:
logger.error(f"保存攻击建议结果到MongoDB失败: {e}")
return result
async def generate_report(self) -> SecurityReportResponse:
"""
......@@ -172,7 +190,7 @@ class SecurityService:
try:
if db.db is not None:
log_entry = result.model_dump()
log_entry["created_at"] = datetime.now(datetime.timezone.utc)
log_entry["created_at"] = datetime.now(timezone.utc)
log_entry["report_date"] = report_data["date"]
log_entry["real_time_data"] = report_data.get("real_time_data", False)
await db.db.security_report_logs.insert_one(log_entry)
......@@ -302,6 +320,10 @@ class SecurityService:
return HistoryQueryResponse(total=total, items=items)
async def _call_llm(self, inputs: Dict[str, Any], model_cls):
if not settings.DIFY_API_URL or not settings.DIFY_API_KEY:
logger.warning("Dify 未配置（缺少 DIFY_API_URL 或 DIFY_API_KEY），跳过 LLM 调用并返回默认结果")
return self._build_default_response(model_cls)
try:
url = f"{settings.DIFY_API_URL.rstrip('/')}/chat-messages"
headers = {
......@@ -327,15 +349,11 @@ class SecurityService:
payload = {
"inputs": inputs,
"query": query_prompt,
"response_mode": "streaming",
"response_mode": settings.DIFY_RESPONSE_MODE,
"conversation_id": "",
"user": "security-system-api",
}
# DEBUG LOG: 打印实际发送的 Payload
logger.info(f"Sending request to Dify. URL: {url}")
logger.info(f"Payload inputs: {json.dumps(inputs, ensure_ascii=False)}")
full_answer = ""
async with httpx.AsyncClient() as client:
async with client.stream("POST", url, json=payload, headers=headers, timeout=120.0) as resp:
......@@ -437,7 +455,7 @@ class SecurityService:
# 但通常思考在前，正文在后。如果没闭合，说明正文还没出来。
# 这里保守处理：如果剩下内容全是思考，那就全删了，返回空串，由上层处理为空的情况。
if '<think>' in text:
text = re.sub(r'<think>.*', '', text, flags=re.DOTALL)
text = re.sub(r'<think>.*', '', text, flags=re.DOTALL)
# 3. 清洗 Markdown 标记
if "```json" in text:
......
microservices/security-service/app/services/zabbix_service.py
View file @ 6e2b6222
......@@ -359,6 +359,14 @@ class ZabbixService:
self.last_sync_time = None
self._initialize_collector()
async def _run_blocking(self, func, *args, **kwargs):
try:
to_thread = asyncio.to_thread
except AttributeError:
loop = asyncio.get_running_loop()
return await loop.run_in_executor(None, lambda: func(*args, **kwargs))
return await to_thread(func, *args, **kwargs)
def _initialize_collector(self):
"""初始化数据采集器"""
try:
......@@ -376,22 +384,22 @@ class ZabbixService:
"""采集设备数据"""
if not self.collector:
raise Exception("Zabbix collector未初始化，请检查Zabbix配置")
return self.collector.get_security_data_for_analysis()
return await self._run_blocking(self.collector.get_security_data_for_analysis)
async def collect_cpu_data(self):
"""采集CPU和硬件数据"""
if not self.collector:
raise Exception("Zabbix collector未初始化，请检查Zabbix配置")
return self.collector.get_cpu_data()
return await self._run_blocking(self.collector.get_cpu_data)
async def collect_network_data(self):
"""采集网络接口数据"""
if not self.collector:
raise Exception("Zabbix collector未初始化，请检查Zabbix配置")
return self.collector.get_network_data()
return await self._run_blocking(self.collector.get_network_data)
async def sync_data(self):
"""同步数据"""
......@@ -426,4 +434,4 @@ class ZabbixService:
return {
"last_sync_time": self.last_sync_time.isoformat() if self.last_sync_time else None,
"collector_initialized": self.collector is not None
}
\ No newline at end of file
}
microservices/security-service/test_refactor.py
View file @ 6e2b6222
#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Security Service 重构测试脚本
测试新的Zabbix集成功能
"""Security Service 本地工具
用途：替代 intelligentPerception/zabbixDataFrom 的采集与联调脚本，支持：
- Zabbix API → 采集 → 生成 JSON 文件
- JSON 文件 → 调用 Security Service REST API
- 保存 Dify 分析结果到文件（服务侧同时落库）
- 查询历史接口并保存结果
"""
import argparse
import asyncio
import sys
import os
import json
import os
import sys
from datetime import datetime
from pathlib import Path
from typing import Any, Dict, Optional
# 添加app目录到Python路径
sys.path.append(os.path.join(os.path.dirname(__file__), 'app'))
import httpx
sys.path.append(os.path.join(os.path.dirname(__file__), "app"))
def _ensure_dir(path: Path) -> Path:
path.mkdir(parents=True, exist_ok=True)
return path
def _write_json(path: Path, data: Any) -> None:
path.write_text(json.dumps(data, indent=2, ensure_ascii=False), encoding="utf-8")
def _read_json(path: Path) -> Any:
return json.loads(path.read_text(encoding="utf-8"))
async def collect_zabbix_json(output_dir: Path) -> Dict[str, Path]:
from services.zabbix_service import ZabbixService
_ensure_dir(output_dir)
zabbix_service = ZabbixService()
status = zabbix_service.get_sync_status()
if not status.get("collector_initialized"):
raise RuntimeError("Zabbix collector 未初始化，请检查 ZABBIX_URL / ZABBIX_USERNAME / ZABBIX_PASSWORD")
print("=" * 70)
print("阶段1：Zabbix 数据采集 → 生成 JSON 文件")
print("=" * 70)
device_data = await zabbix_service.collect_device_data()
cpu_data = await zabbix_service.collect_cpu_data()
network_data = await zabbix_service.collect_network_data()
analysis_input = {"devices": device_data.get("devices", [])}
attack_target = "unknown"
attack_logs = ""
devices = device_data.get("devices", [])
if devices:
picked = next((d for d in devices if d.get("logs")), devices[0])
attack_target = picked.get("name") or "unknown"
logs = picked.get("logs") or []
attack_logs = "\n".join(str(x) for x in logs[:10])
attack_advice_input = {
"attack_type": "Suspicious Activity",
"target_device": attack_target,
"severity": "high",
"logs": attack_logs,
}
paths = {
"zabbix_devices.json": output_dir / "zabbix_devices.json",
"zabbix_cpu.json": output_dir / "zabbix_cpu.json",
"zabbix_network.json": output_dir / "zabbix_network.json",
"analysis_input.json": output_dir / "analysis_input.json",
"attack_advice_input.json": output_dir / "attack_advice_input.json",
}
_write_json(paths["zabbix_devices.json"], device_data)
_write_json(paths["zabbix_cpu.json"], cpu_data)
_write_json(paths["zabbix_network.json"], network_data)
_write_json(paths["analysis_input.json"], analysis_input)
_write_json(paths["attack_advice_input.json"], attack_advice_input)
print(f"✅ 已生成: {paths['analysis_input.json']}")
print(f"✅ 已生成: {paths['attack_advice_input.json']}")
print(f"✅ 设备数量: {len(analysis_input['devices'])}")
return paths
def _auth_headers(token: str) -> Dict[str, str]:
return {"Authorization": f"Bearer {token}", "Content-Type": "application/json"}
async def test_zabbix_service():
"""测试Zabbix服务"""
print("="*70)
print("测试 Zabbix Service")
print("="*70)
try:
from services.zabbix_service import ZabbixService
# 创建Zabbix服务实例
zabbix_service = ZabbixService()
# 测试获取同步状态
print("1. 检查Zabbix服务状态...")
status = zabbix_service.get_sync_status()
print(f" 初始化状态: {'✅ 成功' if status['collector_initialized'] else '❌ 失败'}")
print(f" 最后同步时间: {status['last_sync_time']}")
# 测试同步数据
print("\n2. 测试Zabbix数据同步...")
sync_result = await zabbix_service.sync_data()
print(f" 同步状态: ✅ 成功")
print(f" 同步结果: {json.dumps(sync_result, indent=2, ensure_ascii=False)}")
# 测试获取设备数据
print("\n3. 测试获取设备数据...")
device_data = await zabbix_service.collect_device_data()
devices = device_data.get("devices", [])
print(f" 设备数量: {len(devices)}")
# 显示前3个设备
for i, device in enumerate(devices[:3]):
print(f" 设备 {i+1}: {device['name']} ({device['type']}) - {device['status']}")
print(f" 日志数量: {len(device['logs'])}")
return True
except Exception as e:
print(f"❌ Zabbix Service测试失败: {e}")
return False
async def test_security_service():
"""测试安全服务"""
print("\n" + "="*70)
print("测试 Security Service")
print("="*70)
async def call_security_service_api(base_url: str, token: str, input_dir: Path, output_dir: Path) -> Dict[str, Path]:
_ensure_dir(output_dir)
print("\n" + "=" * 70)
print("阶段2：JSON 文件 → Security Service REST API → 保存结果")
print("=" * 70)
analysis_input = _read_json(input_dir / "analysis_input.json")
attack_input = _read_json(input_dir / "attack_advice_input.json")
endpoints = {
"analysis": f"{base_url.rstrip('/')}/api/v1/security/analysis",
"attack_advice": f"{base_url.rstrip('/')}/api/v1/security/attack-advice",
"monitor": f"{base_url.rstrip('/')}/api/v1/security/monitor",
"report": f"{base_url.rstrip('/')}/api/v1/security/report",
"analysis_history": f"{base_url.rstrip('/')}/api/v1/security/analysis/history?limit=20",
"attack_history": f"{base_url.rstrip('/')}/api/v1/security/attack-advice/history?limit=20",
}
out_paths = {
"analysis_result.json": output_dir / "analysis_result.json",
"attack_advice_result.json": output_dir / "attack_advice_result.json",
"monitor_result.json": output_dir / "monitor_result.json",
"report_result.json": output_dir / "report_result.json",
"analysis_history.json": output_dir / "analysis_history.json",
"attack_advice_history.json": output_dir / "attack_advice_history.json",
}
async with httpx.AsyncClient(timeout=120.0) as client:
r1 = await client.post(endpoints["analysis"], headers=_auth_headers(token), json=analysis_input)
r1.raise_for_status()
analysis_res = r1.json()
_write_json(out_paths["analysis_result.json"], analysis_res)
print(f"✅ /analysis 完成，risk_level={analysis_res.get('risk_level')}")
r2 = await client.post(endpoints["attack_advice"], headers=_auth_headers(token), json=attack_input)
r2.raise_for_status()
attack_res = r2.json()
_write_json(out_paths["attack_advice_result.json"], attack_res)
print(f"✅ /attack-advice 完成，immediate_actions={len(attack_res.get('immediate_actions', []))}")
r3 = await client.get(endpoints["monitor"], headers=_auth_headers(token))
r3.raise_for_status()
monitor_res = r3.json()
_write_json(out_paths["monitor_result.json"], monitor_res)
print(f"✅ /monitor 完成，detected={len(monitor_res.get('detected_vulnerabilities', []))}")
r4 = await client.get(endpoints["report"], headers=_auth_headers(token))
r4.raise_for_status()
report_res = r4.json()
_write_json(out_paths["report_result.json"], report_res)
print(f"✅ /report 完成，date={report_res.get('date')}")
r5 = await client.get(endpoints["analysis_history"], headers=_auth_headers(token))
r5.raise_for_status()
_write_json(out_paths["analysis_history.json"], r5.json())
print("✅ /analysis/history 完成")
r6 = await client.get(endpoints["attack_history"], headers=_auth_headers(token))
r6.raise_for_status()
_write_json(out_paths["attack_advice_history.json"], r6.json())
print("✅ /attack-advice/history 完成")
return out_paths
async def internal_smoke_test() -> bool:
print("=" * 70)
print("阶段0：内部冒烟（不经 REST，仅验证类可用）")
print("=" * 70)
try:
from services.zabbix_service import ZabbixService
from services.analysis import SecurityService
# 创建服务实例
from schemas.payloads import DeviceInfo
zabbix_service = ZabbixService()
status = zabbix_service.get_sync_status()
print(f"Zabbix collector: {'✅ 已初始化' if status.get('collector_initialized') else '❌ 未初始化'}")
security_service = SecurityService(zabbix_service=zabbix_service)
# 测试风险分析
print("1. 测试风险分析功能...")
analysis_result = await security_service.analyze_risks()
print(f" 分析状态: ✅ 成功")
print(f" 风险级别: {analysis_result.risk_level}")
print(f" 漏洞数量: {len(analysis_result.vulnerabilities)}")
print(f" 建议数量: {len(analysis_result.suggestions)}")
# 测试监控功能
print("\n2. 测试风险监控功能...")
monitor_result = await security_service.monitor_risks()
print(f" 监控状态: ✅ 成功")
print(f" 检测漏洞: {len(monitor_result.detected_vulnerabilities)}")
print(f" 合规风险: {len(monitor_result.compliance_risks)}")
# 测试报告生成
print("\n3. 测试报告生成功能...")
report_result = await security_service.generate_report()
print(f" 报告状态: ✅ 成功")
print(f" 报告日期: {report_result.date}")
print(f" 整体状态: {report_result.overall_status}")
# 测试攻击建议
print("\n4. 测试攻击建议功能...")
advice_result = await security_service.get_attack_advice(
attack_type="Port Scan",
target="Web Server",
logs="Port scan detected from external IP"
res = await security_service.analyze_risks(
[
DeviceInfo(
id="smoke-1",
name="Smoke-Device",
type="server",
status="up",
logs=["CPU high", "Multiple failed login attempts"],
)
]
)
print(f" 建议状态: ✅ 成功")
print(f" 即时行动: {len(advice_result.immediate_actions)}")
print(f" 分析结果: {advice_result.analysis[:100]}...")
print(f"Security analyze_risks: ✅ 返回 risk_level={res.risk_level}")
return True
except Exception as e:
print(f"❌ Security Service测试失败: {e}")
print(f"❌ 内部冒烟失败: {e}")
return False
async def main():
"""主测试函数"""
print("Security Service 重构验证测试")
print(f"测试时间: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
print("此脚本测试新的Zabbix集成功能...")
# 测试结果
results = {
"zabbix_service": False,
"security_service": False,
"overall": False
}
# 测试Zabbix服务
results["zabbix_service"] = await test_zabbix_service()
# 测试安全服务
results["security_service"] = await test_security_service()
# 总体结果
results["overall"] = results["zabbix_service"] and results["security_service"]
# 输出测试总结
print("\n" + "="*70)
print("测试总结")
print("="*70)
print(f"Zabbix Service: {'✅ 通过' if results['zabbix_service'] else '❌ 失败'}")
print(f"Security Service: {'✅ 通过' if results['security_service'] else '❌ 失败'}")
print(f"总体结果: {'✅ 全部通过' if results['overall'] else '❌ 有失败项'}")
if results['overall']:
print("\n🎉 重构成功！所有功能正常工作。")
print("现在可以使用以下API端点:")
print(" POST /api/v1/security/zabbix/sync - 手动同步Zabbix数据")
print(" GET /api/v1/security/zabbix/status - 获取Zabbix服务状态")
print(" POST /api/v1/security/zabbix/devices - 获取设备列表")
print(" POST /api/v1/security/analysis - 安全分析")
print(" GET /api/v1/security/monitor - 风险监控")
print(" GET /api/v1/security/report - 生成报告")
else:
print("\n⚠️ 测试失败，请检查:")
print(" 1. Zabbix服务器配置")
print(" 2. 环境变量设置")
print(" 3. 网络连接")
print(" 4. 依赖安装")
return results["overall"]
def _env_or_arg(value: Optional[str], env_key: str) -> Optional[str]:
return value if value else os.getenv(env_key)
async def main() -> int:
parser = argparse.ArgumentParser(prog="test_refactor.py")
sub = parser.add_subparsers(dest="cmd", required=False)
p_collect = sub.add_parser("collect", help="采集Zabbix并生成JSON文件")
p_collect.add_argument("--out", default=".\\out", help="输出目录")
p_call = sub.add_parser("call-api", help="读取JSON并调用Security Service REST API")
p_call.add_argument("--base-url", default=None, help="服务地址，例如 http://localhost:8003 或 http://localhost:8080")
p_call.add_argument("--token", default=None, help="管理员JWT，或使用环境变量 ADMIN_TOKEN")
p_call.add_argument("--in", dest="in_dir", default=".\\out", help="输入目录（包含 analysis_input.json 等）")
p_call.add_argument("--out", dest="out_dir", default=".\\out", help="输出目录")
p_all = sub.add_parser("run-all", help="采集→生成JSON→调用REST→保存结果→查询历史")
p_all.add_argument("--base-url", default=None, help="服务地址，例如 http://localhost:8003 或 http://localhost:8080")
p_all.add_argument("--token", default=None, help="管理员JWT，或使用环境变量 ADMIN_TOKEN")
p_all.add_argument("--out", default=".\\out", help="输出目录")
sub.add_parser("smoke", help="内部冒烟测试")
args = parser.parse_args()
print("Security Service 一体化工具")
print(f"时间: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
if args.cmd in (None, "smoke"):
ok = await internal_smoke_test()
return 0 if ok else 1
if args.cmd == "collect":
out_dir = Path(args.out).resolve()
await collect_zabbix_json(out_dir)
return 0
if args.cmd == "call-api":
base_url = _env_or_arg(args.base_url, "SECURITY_SERVICE_URL") or "http://localhost:8003"
token = _env_or_arg(args.token, "ADMIN_TOKEN")
if not token:
raise RuntimeError("缺少管理员JWT：请传 --token 或设置环境变量 ADMIN_TOKEN")
in_dir = Path(args.in_dir).resolve()
out_dir = Path(args.out_dir).resolve()
await call_security_service_api(base_url, token, in_dir, out_dir)
return 0
if args.cmd == "run-all":
base_url = _env_or_arg(args.base_url, "SECURITY_SERVICE_URL") or "http://localhost:8003"
token = _env_or_arg(args.token, "ADMIN_TOKEN")
if not token:
raise RuntimeError("缺少管理员JWT：请传 --token 或设置环境变量 ADMIN_TOKEN")
out_dir = Path(args.out).resolve()
await collect_zabbix_json(out_dir)
await call_security_service_api(base_url, token, out_dir, out_dir)
return 0
raise RuntimeError(f"未知命令: {args.cmd}")
if __name__ == "__main__":
# 运行测试
success = asyncio.run(main())
sys.exit(0 if success else 1)
\ No newline at end of file
try:
sys.exit(asyncio.run(main()))
except KeyboardInterrupt:
sys.exit(130)
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment