实时数仓推动下漏洞管理的技术演进方向与关键挑战

RTVulnFlow/
├── config.yaml                     # 项目配置文件
├── requirements.txt                # Python依赖
├── run.py                         # 项目主启动脚本
├── src/
│   ├── __init__.py
│   ├── data_simulator.py          # 模拟数据源（资产、漏洞、性能数据）
│   ├── stream_processor.py        # Flink流处理作业核心逻辑
│   ├── vuln_detector.py           # 漏洞检测逻辑（规则+轻量ML）
│   ├── repair_orchestrator.py     # 修复联动器
│   └── metrics_monitor.py         # 简易监控器
└── tests/
    └── test_detector.py           # 单元测试示例

# RTVulnFlow 核心配置
kafka:
  bootstrap_servers: "localhost:9092"
  topics:
    asset_events: "rtvuln.asset.events"
    vuln_raw_events: "rtvuln.vuln.raw"
    vuln_enriched_events: "rtvuln.vuln.enriched"
    performance_metrics: "rtvuln.perf.metrics"

flink:
  checkpoint_interval: 30000 # ms
  parallelism: 2

detection:
  rule_engine:
    critical_cves: ["CVE-2021-44228", "CVE-2021-45046"] # 严重CVE名单
    severity_threshold: "HIGH"
  ml_model:
    # 示例：使用本地模型文件或在线API端点
    model_path: "models/vuln_risk_predictor.pkl"
    enabled: true

remediation:
  webhook_url: "https://hook.example.com/security-remediation"
  severity_trigger: "CRITICAL"
  enabled: true

monitoring:
  metric_export_interval: 60000 # ms

apache-flink==1.16.1
pyyaml==6.0
requests==2.31.0
scikit-learn==1.3.0
kafka-python==2.0.2
# 用于模拟数据生成
faker==19.6.2

import json
import time
import random
from datetime import datetime
from dataclasses import dataclass, asdict
from enum import Enum
import yaml
from kafka import KafkaProducer
from faker import Faker

fake = Faker()

class AssetType(Enum):
    CONTAINER = "container"
    VM = "virtual_machine"
    K8S_POD = "kubernetes_pod"
    RDS = "rds_instance"

class VulnerabilitySeverity(Enum):
    LOW = "LOW"
    MEDIUM = "MEDIUM"
    HIGH = "HIGH"
    CRITICAL = "CRITICAL"

@dataclass
class AssetEvent:
    asset_id: str
    asset_type: AssetType
    hostname: str
    ip_address: str
    owner: str
    department: str
    criticality: int  # 1-5, 5为最核心
    timestamp: str

@dataclass
class VulnerabilityEvent:
    scan_id: str
    asset_id: str
    cve_id: str
    severity: VulnerabilitySeverity
    description: str
    cvss_score: float
    package_name: str
    package_version: str
    timestamp: str

@dataclass
class PerformanceMetric:
    asset_id: str
    cpu_utilization: float  # 百分比
    memory_utilization: float
    network_in: float  # MB/s
    network_out: float
    timestamp: str

class DataSimulator:
    """
    模拟生成资产、漏洞和性能指标数据，并发送到Kafka。
    这是一个简化版，实际项目应从真实扫描器、CMDB等获取数据。
    """
    def __init__(self, config_path='config.yaml'):
        with open(config_path, 'r') as f:
            self.config = yaml.safe_load(f)
        kafka_config = self.config['kafka']
        self.producer = KafkaProducer(
            bootstrap_servers=kafka_config['bootstrap_servers'].split(','),
            value_serializer=lambda v: json.dumps(v).encode('utf-8')
        )
        self.topics = kafka_config['topics']
        self.known_assets = []

    def _generate_asset(self):
        """生成一个模拟资产"""
        asset_types = list(AssetType)
        asset = AssetEvent(
            asset_id=fake.uuid4(),
            asset_type=random.choice(asset_types),
            hostname=fake.hostname(),
            ip_address=fake.ipv4(),
            owner=fake.name(),
            department=fake.job(),
            criticality=random.randint(1,5),
            timestamp=datetime.utcnow().isoformat() + 'Z'
        )
        self.known_assets.append(asset.asset_id)
        return asset

    def _generate_vuln_for_asset(self, asset_id):
        """为指定资产生成一个模拟漏洞事件"""
        cve_list = ["CVE-2021-44228", "CVE-2021-45046", "CVE-2022-22963", "CVE-2022-42889", "CVE-2023-12345"]
        severities = list(VulnerabilitySeverity)
        # 使CRITICAL和HIGH出现的概率更低
        severity = random.choices(severities, weights=[0.3, 0.4, 0.2, 0.1])[0]
        return VulnerabilityEvent(
            scan_id=fake.uuid4(),
            asset_id=asset_id,
            cve_id=random.choice(cve_list),
            severity=severity,
            description=fake.sentence(),
            cvss_score=round(random.uniform(2.0, 10.0), 1),
            package_name=fake.word(),
            package_version=f"{random.randint(1,5)}.{random.randint(0,9)}.{random.randint(0,9)}",
            timestamp=datetime.utcnow().isoformat() + 'Z'
        )

    def _generate_perf_for_asset(self, asset_id):
        """为指定资产生成模拟性能指标"""
        return PerformanceMetric(
            asset_id=asset_id,
            cpu_utilization=round(random.uniform(5.0, 95.0), 2),
            memory_utilization=round(random.uniform(10.0, 99.0), 2),
            network_in=round(random.uniform(1.0, 100.0), 2),
            network_out=round(random.uniform(1.0, 100.0), 2),
            timestamp=datetime.utcnow().isoformat() + 'Z'
        )

    def run(self, duration_seconds=60, interval=2):
        """
        运行模拟器，持续生成数据。
        Args:
            duration_seconds: 总运行时间（秒）
            interval: 每次发送的间隔（秒）
        """
        print(f"Starting data simulator for {duration_seconds} seconds...")
        end_time = time.time() + duration_seconds
        asset_creation_counter = 0

        while time.time() < end_time:
            # 1. 偶尔创建新资产
            if random.random() < 0.1 or asset_creation_counter < 5:
                new_asset = self._generate_asset()
                self.producer.send(self.topics['asset_events'], asdict(new_asset))
                print(f"[Asset Event] Created: {new_asset.hostname}({new_asset.asset_id})")
                asset_creation_counter += 1

            # 2. 为随机已有资产生成漏洞
            if self.known_assets and random.random() < 0.7:
                target_asset = random.choice(self.known_assets)
                vuln = self._generate_vuln_for_asset(target_asset)
                self.producer.send(self.topics['vuln_raw_events'], asdict(vuln))
                print(f"[Vuln Event] Found: {vuln.cve_id}({vuln.severity.value}) on {target_asset}")

            # 3. 为随机已有资产生成性能指标
            if self.known_assets:
                target_asset = random.choice(self.known_assets)
                perf = self._generate_perf_for_asset(target_asset)
                self.producer.send(self.topics['performance_metrics'], asdict(perf))

            time.sleep(interval)
        print("Data simulation finished.")
        self.producer.flush()
        self.producer.close()

if __name__ == "__main__":
    simulator = DataSimulator()
    simulator.run(duration_seconds=30, interval=1)

import pickle
import numpy as np
from enum import Enum
from dataclasses import dataclass
from typing import Optional

@dataclass
class EnrichedVulnerability:
    """富化后的漏洞事件，包含原始信息及分析结果"""
    original_event: dict
    risk_score: float  # 0-100，综合风险评分
    requires_immediate_action: bool
    business_unit: str
    performance_impact: float  # 预估性能影响系数 0-1
    recommended_action: str

class RiskPredictor:
    """
    轻量级机器学习风险预测器。
    示例：使用简单的特征（CVSS分数，资产关键性，性能压力）预测风险评分。
    生产环境应使用更复杂的模型和特征工程。
    """
    def __init__(self, model_path: Optional[str] = None):
        if model_path:
            # 加载预训练模型
            with open(model_path, 'rb') as f:
                self.model = pickle.load(f)
        else:
            # 创建一个虚拟模型（实际应训练）
            from sklearn.linear_model import LinearRegression
            self.model = LinearRegression()
            # 用虚拟数据拟合一下，避免predict时报错
            dummy_X = np.array([[5.0, 3, 0.5], [8.0, 5, 0.9]])
            dummy_y = np.array([40, 85])
            self.model.fit(dummy_X, dummy_y)

    def predict(self, cvss_score: float, asset_criticality: int, perf_pressure: float) -> float:
        """预测风险分数"""
        features = np.array([[cvss_score, asset_criticality, perf_pressure]])
        score = self.model.predict(features)[0]
        return np.clip(score, 0, 100)  # 限制在0-100

class VulnDetector:
    """
    漏洞检测与富化引擎。
    结合规则引擎（硬编码规则）与ML模型进行实时风险评估。
    """
    def __init__(self, config: dict, risk_predictor: RiskPredictor):
        self.rule_config = config['detection']['rule_engine']
        self.ml_enabled = config['detection']['ml_model']['enabled']
        self.risk_predictor = risk_predictor
        # 模拟资产上下文缓存 (生产环境应来自实时数仓的维表或实时查询)
        self.asset_context_cache = {}

    def update_asset_context(self, asset_id: str, context: dict):
        """更新资产上下文信息（可从资产事件流中更新）"""
        self.asset_context_cache[asset_id] = context

    def calculate_performance_impact(self, perf_metric: dict) -> float:
        """根据性能指标计算性能影响系数"""
        cpu = perf_metric.get('cpu_utilization', 50) / 100.0
        mem = perf_metric.get('memory_utilization', 50) / 100.0
        # 简单加权平均，可根据实际情况调整
        impact = (cpu * 0.6 + mem * 0.4)
        return min(impact, 1.0)

    def detect_and_enrich(self, vuln_event: dict, perf_metric: dict = None) -> EnrichedVulnerability:
        """
        核心检测与富化函数。
        Args:
            vuln_event: 原始漏洞事件字典
            perf_metric: 可选的关联性能指标
        Returns:
            EnrichedVulnerability
        """
        asset_id = vuln_event['asset_id']
        cve_id = vuln_event['cve_id']
        severity = vuln_event['severity']
        cvss_score = vuln_event['cvss_score']

        # 1. 获取资产上下文
        asset_context = self.asset_context_cache.get(asset_id, {})
        asset_criticality = asset_context.get('criticality', 3)  # 默认值
        business_unit = asset_context.get('department', 'Unknown')

        # 2. 计算性能影响
        perf_impact = self.calculate_performance_impact(perf_metric) if perf_metric else 0.3

        # 3. 规则匹配（硬编码规则示例）
        requires_immediate_action = False
        if cve_id in self.rule_config['critical_cves']:
            requires_immediate_action = True
        if severity == 'CRITICAL' or (severity == 'HIGH' and asset_criticality >= 4):
            requires_immediate_action = True

        # 4. ML风险评分
        risk_score = 0.0
        if self.ml_enabled:
            risk_score = self.risk_predictor.predict(cvss_score, asset_criticality, perf_impact)
        else:
            # 基于规则的简单评分
            risk_score = cvss_score * 10 * (asset_criticality / 5.0) * (0.5 + 0.5 * perf_impact)

        # 5. 生成建议动作
        recommended_action = "Monitor"
        if risk_score > 70:
            recommended_action = "Patch within 24 hours"
        if requires_immediate_action:
            recommended_action = "Emergency patch or isolate asset"

        return EnrichedVulnerability(
            original_event=vuln_event,
            risk_score=round(risk_score, 2),
            requires_immediate_action=requires_immediate_action,
            business_unit=business_unit,
            performance_impact=round(perf_impact, 2),
            recommended_action=recommended_action
        )

import json
import yaml
from pyflink.datastream import StreamExecutionEnvironment, TimeCharacteristic
from pyflink.datastream.connectors import FlinkKafkaConsumer, FlinkKafkaProducer
from pyflink.datastream.formats import JsonRowDeserializationSchema, JsonRowSerializationSchema
from pyflink.common.serialization import SimpleStringSchema
from pyflink.common.typeinfo import Types
from pyflink.datastream.functions import ProcessFunction, RuntimeContext
from pyflink.common import WatermarkStrategy, Time
from pyflink.datastream.window import TumblingEventTimeWindows
from pyflink.datastream.state import ValueStateDescriptor, ValueState
import logging

from src.vuln_detector import VulnDetector, RiskPredictor

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

class AssetContextProcessFunction(ProcessFunction):
    """处理资产事件，用于更新资产上下文状态"""
    def __init__(self):
        self.asset_state = None

    def open(self, runtime_context: RuntimeContext):
        state_descriptor = ValueStateDescriptor("asset_context", Types.PY_DICT)
        self.asset_state = runtime_context.get_state(state_descriptor)

    def process_element(self, value, ctx):
        # value 是一个包含资产事件数据的元组或字典 (由上游DataStream类型决定)
        # 简化处理：我们假设value已经是dict
        asset_id = value['asset_id']
        # 将资产信息存入状态，供漏洞事件关联时查询
        ctx.output(self.tag, (asset_id, value)) # 输出到侧输出流
        # 也可以更新一个全局状态（这里简化，实际用BroadcastState或KeyedState）
        logger.debug(f"Processed asset event for {asset_id}")
        # 注意：这里简化了状态管理。生产环境需考虑状态TTL和规模。

class VulnEnrichmentProcessFunction(ProcessFunction):
    """核心的漏洞富化处理函数"""
    def __init__(self, config):
        self.config = config
        self.detector = None

    def open(self, runtime_context: RuntimeContext):
        # 初始化检测器（每个并行子任务一个实例）
        predictor = RiskPredictor(self.config['detection']['ml_model'].get('model_path'))
        self.detector = VulnDetector(self.config, predictor)

    def process_element(self, vuln_event, ctx):
        """
        vuln_event: 从Kafka消费的原始漏洞事件 (dict)
        此函数应关联资产上下文和最近的性能指标，然后调用检测器。
        简化版：我们假设资产上下文和性能指标已通过广播流或连接操作关联好了。
        本例中，我们模拟从状态获取。
        """
        asset_id = vuln_event['asset_id']
        # TODO: 实际应从KeyedState中获取该asset_id对应的上下文和性能指标
        # 这里为了示例可运行，我们模拟一些数据
        simulated_asset_context = {'criticality': 4, 'department': 'Engineering'}
        simulated_perf_metric = {'cpu_utilization': 80.0, 'memory_utilization': 65.0}

        self.detector.update_asset_context(asset_id, simulated_asset_context)
        enriched_vuln = self.detector.detect_and_enrich(vuln_event, simulated_perf_metric)

        # 转换为可序列化的字典
        output_event = {
            'enriched_timestamp': ctx.timestamp(), # 处理时间
            'asset_id': asset_id,
            'cve_id': vuln_event['cve_id'],
            'original_severity': vuln_event['severity'],
            'risk_score': enriched_vuln.risk_score,
            'business_unit': enriched_vuln.business_unit,
            'performance_impact': enriched_vuln.performance_impact,
            'requires_immediate_action': enriched_vuln.requires_immediate_action,
            'recommended_action': enriched_vuln.recommended_action,
            'original_event_summary': f"{vuln_event['cve_id']} on {asset_id}"
        }
        logger.info(f"Enriched Vuln: {output_event['cve_id']} -> Risk: {output_event['risk_score']}, Action: {output_event['recommended_action']}")
        return output_event

def create_kafka_source(env, topic, bootstrap_servers, group_id='rtvuln-group'):
    """创建Kafka源"""
    properties = {
        'bootstrap.servers': bootstrap_servers,
        'group.id': group_id,
        'auto.offset.reset': 'latest'
    }
    # 使用简单的String反序列化，然后在后续算子中解析JSON
    kafka_consumer = FlinkKafkaConsumer(
        topics=topic,
        deserialization_schema=SimpleStringSchema(),
        properties=properties
    )
    # 设置从最新开始消费，适合测试
    kafka_consumer.set_start_from_latest()
    return env.add_source(kafka_consumer)

def create_kafka_sink(topic, bootstrap_servers):
    """创建Kafka sink"""
    properties = {
        'bootstrap.servers': bootstrap_servers
    }
    # 序列化为JSON字符串
    serialization_schema = JsonRowSerializationSchema.Builder() \
        .with_type_info(Types.ROW([
            Types.STRING(), # enriched_timestamp
            Types.STRING(), # asset_id
            Types.STRING(), # cve_id
            Types.STRING(), # original_severity
            Types.DOUBLE(), # risk_score
            Types.STRING(), # business_unit
            Types.DOUBLE(), # performance_impact
            Types.BOOLEAN(), # requires_immediate_action
            Types.STRING(), # recommended_action
            Types.STRING()  # original_event_summary
        ])).build()

    kafka_producer = FlinkKafkaProducer(
        topic=topic,
        serialization_schema=serialization_schema,
        producer_config=properties
    )
    return kafka_producer

def main_job(config_path='config.yaml'):
    # 1. 加载配置
    with open(config_path, 'r') as f:
        config = yaml.safe_load(f)
    kafka_config = config['kafka']

    # 2. 设置流执行环境
    env = StreamExecutionEnvironment.get_execution_environment()
    env.set_stream_time_characteristic(TimeCharacteristic.EventTime) # 或ProcessingTime
    env.set_parallelism(config['flink']['parallelism'])
    # 添加Flink Kafka连接器JAR（本地运行需指定）
    # env.add_jars("file:///path/to/flink-sql-connector-kafka.jar")

    # 3. 创建数据源
    vuln_source = create_kafka_source(
        env,
        kafka_config['topics']['vuln_raw_events'],
        kafka_config['bootstrap_servers']
    ).name("Vulnerability Source")

    # 4. 数据转换与处理
    # 4.1 解析JSON字符串
    parsed_stream = vuln_source.map(
        lambda s: json.loads(s),
        output_type=Types.PY_DICT
    ).name("Parse JSON")

    # 4.2 关键处理：富化漏洞
    enriched_stream = parsed_stream.map(
        lambda e: VulnEnrichmentProcessFunction(config).process_element(e, None), # 简化，未传递Context
        output_type=Types.PY_DICT
    ).name("Vulnerability Enrichment")

    # 5. 创建并添加Sink
    sink = create_kafka_sink(
        kafka_config['topics']['vuln_enriched_events'],
        kafka_config['bootstrap_servers']
    ).name("Enriched Events Sink")
    # 需要将Python字典转换为Row对象
    from pyflink.common import Row
    enriched_stream.map(
        lambda d: Row(
            d.get('enriched_timestamp', ''),
            d['asset_id'],
            d['cve_id'],
            d['original_severity'],
            d['risk_score'],
            d['business_unit'],
            d['performance_impact'],
            d['requires_immediate_action'],
            d['recommended_action'],
            d['original_event_summary']
        ),
        output_type=Types.ROW([
            Types.STRING(), Types.STRING(), Types.STRING(), Types.STRING(),
            Types.DOUBLE(), Types.STRING(), Types.DOUBLE(), Types.BOOLEAN(),
            Types.STRING(), Types.STRING()
        ])
    ).name("To Row").add_sink(sink)

    # 6. 执行作业
    logger.info("Submitting Flink job: Real-Time Vulnerability Enrichment Pipeline")
    env.execute("RTVulnFlow-Enrichment")

if __name__ == "__main__":
    main_job()

import requests
import json
import logging
from dataclasses import asdict
from typing import Dict, Any

logger = logging.getLogger(__name__)

class RepairOrchestrator:
    """
    修复联动器。
    接收富化后的高风险漏洞事件，并通过Webhook、API调用等方式触发外部系统进行修复。
    """
    def __init__(self, config: Dict[str, Any]):
        self.config = config['remediation']
        self.webhook_url = self.config.get('webhook_url')
        self.severity_trigger = self.config.get('severity_trigger', 'CRITICAL')
        self.enabled = self.config.get('enabled', False)

    def should_trigger(self, enriched_event: Dict[str, Any]) -> bool:
        """判断是否需要触发修复流程"""
        if not self.enabled:
            return False
        immediate = enriched_event.get('requires_immediate_action', False)
        risk_score = enriched_event.get('risk_score', 0)
        # 策略：需要立即行动 或 风险分数大于阈值
        return immediate or risk_score > 75

    def trigger_remediation(self, enriched_event: Dict[str, Any]) -> bool:
        """
        触发修复动作。
        返回是否成功。
        """
        if not self.should_trigger(enriched_event):
            logger.info(f"No trigger for {enriched_event.get('cve_id')}. Risk score: {enriched_event.get('risk_score')}")
            return False

        payload = self._build_payload(enriched_event)
        logger.warning(f"ATTEMPTING REMEDIATION TRIGGER for {enriched_event.get('cve_id')} on {enriched_event.get('asset_id')}")
        logger.info(f"Payload: {json.dumps(payload, indent=2)}")

        # 在实际项目中，这里可能是调用Jira、ServiceNow、Ansible Tower等的API
        if self.webhook_url and self.webhook_url.startswith('http'):
            try:
                response = requests.post(
                    self.webhook_url,
                    json=payload,
                    headers={'Content-Type': 'application/json'},
                    timeout=10
                )
                if response.status_code in [200, 201, 202]:
                    logger.info(f"Remediation triggered successfully via webhook. Response: {response.status_code}")
                    return True
                else:
                    logger.error(f"Webhook call failed. Status: {response.status_code}, Body: {response.text}")
                    return False
            except Exception as e:
                logger.exception(f"Exception while calling remediation webhook: {e}")
                return False
        else:
            # 模拟成功，仅记录
            logger.info(f"[SIMULATION] Remediation would be triggered for {payload['issue_key']}")
            return True

    def _build_payload(self, event: Dict[str, Any]) -> Dict[str, Any]:
        """构建发送给外部系统的负载"""
        # 根据不同的下游系统，构建不同的payload结构
        # 这里是一个通用示例
        return {
            "issue_type": "Security Patch",
            "issue_key": f"VULN-{event['cve_id']}-{event['asset_id'][:8]}",
            "title": f"Critical Vulnerability {event['cve_id']} detected on {event['asset_id']}",
            "description": f"Risk Score: {event['risk_score']}. Impact on {event['business_unit']}. Performance impact factor: {event['performance_impact']}. Action required: {event['recommended_action']}",
            "priority": "P1" if event['requires_immediate_action'] else "P2",
            "asset_id": event['asset_id'],
            "cve_id": event['cve_id'],
            "risk_score": event['risk_score'],
            "recommended_action": event['recommended_action'],
            "source": "RTVulnFlow"
        }

import time
import threading
from collections import defaultdict
from typing import Dict, List
import logging

logger = logging.getLogger(__name__)

class SimpleMetricsMonitor:
    """
    简易内存指标监控器。
    生产环境应使用Prometheus, InfluxDB等。
    本类用于演示如何收集和暴露处理管道的内部指标。
    """
    def __init__(self, export_interval_seconds: int = 60):
        self.export_interval = export_interval_seconds
        self.metrics = defaultdict(int)
        self.lock = threading.Lock()
        self.running = False
        self.thread = None

    def increment_counter(self, metric_name: str, value: int = 1):
        """递增计数器"""
        with self.lock:
            self.metrics[metric_name] += value

    def set_gauge(self, metric_name: str, value: float):
        """设置仪表盘值（记录最新值）"""
        with self.lock:
            self.metrics[metric_name] = value

    def record_timer(self, metric_name: str, duration_ms: float):
        """记录耗时（简单平均）"""
        count_metric = f"{metric_name}_count"
        sum_metric = f"{metric_name}_sum_ms"
        with self.lock:
            self.metrics[count_metric] += 1
            self.metrics[sum_metric] += duration_ms
            # 计算平均时间
            avg_metric = f"{metric_name}_avg_ms"
            count = self.metrics[count_metric]
            if count > 0:
                self.metrics[avg_metric] = self.metrics[sum_metric] / count

    def get_metrics_snapshot(self) -> Dict[str, float]:
        """获取当前指标快照"""
        with self.lock:
            return dict(self.metrics)

    def _export_loop(self):
        """定期导出指标到日志（或外部系统）"""
        while self.running:
            time.sleep(self.export_interval)
            snapshot = self.get_metrics_snapshot()
            if snapshot:
                logger.info("[METRICS MONITOR] ====== Snapshot ======")
                for key, value in snapshot.items():
                    logger.info(f"  {key}: {value}")
                logger.info("[METRICS MONITOR] ======================")

    def start(self):
        """启动后台导出线程"""
        if not self.running:
            self.running = True
            self.thread = threading.Thread(target=self._export_loop, daemon=True)
            self.thread.start()
            logger.info("Metrics monitor started.")

    def stop(self):
        """停止监控器"""
        self.running = False
        if self.thread:
            self.thread.join(timeout=5)
        logger.info("Metrics monitor stopped.")

#!/usr/bin/env python3
"""
RTVulnFlow 主启动脚本。
根据参数启动模拟器、流处理作业或修复联动消费端。
"""
import argparse
import sys
import os
from src.data_simulator import DataSimulator
from src.stream_processor import main_job as start_flink_job
from src.repair_orchestrator import RepairOrchestrator
from src.metrics_monitor import SimpleMetricsMonitor
import yaml
import logging
import time
from kafka import KafkaConsumer
import json

logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

def start_simulator(config_path, duration, interval):
    """启动数据模拟器"""
    logger.info("Starting Data Simulator...")
    simulator = DataSimulator(config_path)
    try:
        simulator.run(duration_seconds=duration, interval=interval)
    except KeyboardInterrupt:
        logger.info("Simulator stopped by user.")

def start_enrichment_pipeline(config_path):
    """启动Flink流处理作业（这里会提交到Flink集群/本地环境）"""
    logger.info("Starting Flink Enrichment Pipeline...")
    # 注意：在本地运行模式下，这通常会在进程内启动一个mini-cluster。
    # 对于生产部署，应通过`flink run`提交JAR包。
    try:
        start_flink_job(config_path)
    except Exception as e:
        logger.exception(f"Flink job failed: {e}")

def start_remediation_consumer(config_path):
    """启动修复联动消费者（从Kafka读取富化事件并触发修复）"""
    logger.info("Starting Remediation Consumer...")
    with open(config_path, 'r') as f:
        config = yaml.safe_load(f)
    kafka_config = config['kafka']
    orchestrator = RepairOrchestrator(config)

    consumer = KafkaConsumer(
        kafka_config['topics']['vuln_enriched_events'],
        bootstrap_servers=kafka_config['bootstrap_servers'].split(','),
        auto_offset_reset='latest',
        enable_auto_commit=True,
        group_id='rtvuln-remediation-group',
        value_deserializer=lambda x: json.loads(x.decode('utf-8'))
    )

    monitor = SimpleMetricsMonitor(export_interval_seconds=30)
    monitor.start()

    try:
        for message in consumer:
            enriched_event = message.value
            logger.debug(f"Received enriched event: {enriched_event['cve_id']}")
            monitor.increment_counter('vuln_enriched_received')
            start_time = time.time()
            success = orchestrator.trigger_remediation(enriched_event)
            duration_ms = (time.time() - start_time) * 1000
            monitor.record_timer('remediation_trigger_latency', duration_ms)
            if success:
                monitor.increment_counter('remediation_triggered_success')
            else:
                monitor.increment_counter('remediation_triggered_skipped')
    except KeyboardInterrupt:
        logger.info("Remediation consumer stopped by user.")
    finally:
        monitor.stop()
        consumer.close()

def main():
    parser = argparse.ArgumentParser(description='RTVulnFlow - Real-Time Vulnerability Management Pipeline')
    parser.add_argument('--config', default='config.yaml', help='Path to configuration file')
    subparsers = parser.add_subparsers(dest='command', help='Available commands')

    parser_sim = subparsers.add_parser('simulate', help='Run data simulator')
    parser_sim.add_argument('--duration', type=int, default=120, help='Simulation duration in seconds')
    parser_sim.add_argument('--interval', type=float, default=2.0, help='Interval between events in seconds')

    subparsers.add_parser('enrich', help='Start the Flink enrichment pipeline (blocks)')

    subparsers.add_parser('remediate', help='Start the remediation consumer (blocks)')

    subparsers.add_parser('all', help='Run simulator, enrichment, and remediation in separate processes (demo)')
    parser_all = subparsers.get_parser('all')
    # 对于'all'命令，通常需要更复杂的进程管理，这里仅作示意。

    args = parser.parse_args()

    if not os.path.exists(args.config):
        logger.error(f"Config file not found: {args.config}")
        sys.exit(1)

    if args.command == 'simulate':
        start_simulator(args.config, args.duration, args.interval)
    elif args.command == 'enrich':
        start_enrichment_pipeline(args.config)
    elif args.command == 'remediate':
        start_remediation_consumer(args.config)
    elif args.command == 'all':
        logger.warning("The 'all' command is a complex demo. In practice, run components separately.")
        logger.info("Suggested manual demo flow:")
        print("""

        1. In Terminal 1: python run.py simulate --duration 300
        2. In Terminal 2: python run.py enrich  (需要Flink环境，可能需要调整代码以在本地运行)
        3. In Terminal 3: python run.py remediate
        """)
    else:
        parser.print_help()

if __name__ == "__main__":
    main()

import sys
import os
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))

from src.vuln_detector import VulnDetector, RiskPredictor
import yaml

def test_detector():
    """测试漏洞检测器基本逻辑"""
    with open('config.yaml', 'r') as f:
        config = yaml.safe_load(f)

    predictor = RiskPredictor() # 使用虚拟模型
    detector = VulnDetector(config, predictor)

    # 模拟一个资产上下文
    detector.update_asset_context('asset-123', {'criticality': 5, 'department': 'Finance'})

    # 模拟一个漏洞事件
    mock_vuln_event = {
        'asset_id': 'asset-123',
        'cve_id': 'CVE-2021-44228', # 在严重名单中
        'severity': 'CRITICAL',
        'cvss_score': 9.8,
        'package_name': 'log4j',
        'package_version': '2.14.0',
        'timestamp': '2023-10-01T12:00:00Z'
    }
    mock_perf_metric = {'cpu_utilization': 90.0, 'memory_utilization': 80.0}

    enriched = detector.detect_and_enrich(mock_vuln_event, mock_perf_metric)

    print(f"Test Result for {mock_vuln_event['cve_id']}:")
    print(f"  Risk Score: {enriched.risk_score}")
    print(f"  Immediate Action: {enriched.requires_immediate_action}")
    print(f"  Recommended Action: {enriched.recommended_action}")
    print(f"  Business Unit: {enriched.business_unit}")
    print(f"  Performance Impact: {enriched.performance_impact}")

    # 简单断言
    assert enriched.requires_immediate_action == True
    assert enriched.recommended_action == "Emergency patch or isolate asset"
    assert enriched.business_unit == "Finance"
    print("\nAll basic assertions passed!")

if __name__ == "__main__":
    test_detector()

# 在项目根目录 RTVulnFlow/
pip install -r requirements.txt

python run.py simulate --duration 300 --interval 1.5

# simple_enricher.py - 简化版消费者，用于演示核心处理逻辑
import json, yaml, time, logging
from kafka import KafkaConsumer, KafkaProducer
from src.vuln_detector import VulnDetector, RiskPredictor

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

def simple_enricher_loop(config_path='config.yaml'):
    with open(config_path) as f:
        config = yaml.safe_load(f)
    kafka_conf = config['kafka']
    predictor = RiskPredictor()
    detector = VulnDetector(config, predictor)

    consumer = KafkaConsumer(
        kafka_conf['topics']['vuln_raw_events'],
        bootstrap_servers=kafka_conf['bootstrap_servers'].split(','),
        auto_offset_reset='latest',
        group_id='rtvuln-simple-enricher',
        value_deserializer=lambda x: json.loads(x.decode('utf-8'))
    )
    producer = KafkaProducer(
        bootstrap_servers=kafka_conf['bootstrap_servers'].split(','),
        value_serializer=lambda v: json.dumps(v).encode('utf-8')
    )

    logger.info("Simple Enricher started. Waiting for raw vuln events...")
    try:
        for message in consumer:
            raw_event = message.value
            logger.info(f"Processing: {raw_event['cve_id']} for asset {raw_event['asset_id']}")
            # 模拟资产上下文和性能指标（实际应从其他topic join）
            asset_context = {'criticality': random.randint(1,5), 'department': 'SimDept'}
            perf_metric = {'cpu_utilization': random.uniform(10,90), 'memory_utilization': random.uniform(20,80)}
            detector.update_asset_context(raw_event['asset_id'], asset_context)
            enriched = detector.detect_and_enrich(raw_event, perf_metric)
            output = {
                'asset_id': enriched.original_event['asset_id'],
                'cve_id': enriched.original_event['cve_id'],
                'risk_score': enriched.risk_score,
                'action_required': enriched.requires_immediate_action,
                'recommendation': enriched.recommended_action,
                'processing_time': time.time()
            }
            producer.send(kafka_conf['topics']['vuln_enriched_events'], output)
            logger.info(f"Enriched and sent: Risk={enriched.risk_score}, Action={enriched.recommended_action}")
    except KeyboardInterrupt:
        logger.info("Simple enricher stopped.")
    finally:
        consumer.close()
        producer.flush()
        producer.close()

if __name__ == "__main__":
    import random
    simple_enricher_loop()

python simple_enricher.py

python run.py remediate

python tests/test_detector.py