AI辅助安全渗透测试_03

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import networkx as nx
import random
import torch
import torch.nn as nn
import torch.optim as optim
from sklearn.ensemble import RandomForestClassifier
from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score, confusion_matrix
from sklearn.preprocessing import StandardScaler
from collections import defaultdict, deque
import pickle
import os
import time
from datetime import datetime
import requests
from bs4 import BeautifulSoup
import json
import subprocess
from tqdm import tqdm
import warnings
warnings.filterwarnings('ignore')

# 设置中文显示
plt.rcParams["font.family"] = ["SimHei", "WenQuanYi Micro Hei", "Heiti TC"]
plt.rcParams['axes.unicode_minus'] = False

class AISecurityTester:
    def __init__(self, model_path=None):
        """初始化AI安全测试器"""
        self.vulnerability_model = None
        self.attack_graph = nx.DiGraph()
        self.knowledge_graph = nx.DiGraph()
        self.asset_db = {}
        self.scan_results = {}
        self.attack_paths = []
        self.vulnerability_db = self._load_vulnerability_db()
        self.scaler = StandardScaler()

        # 如果提供了模型路径，则加载预训练模型
        if model_path and os.path.exists(model_path):
            self._load_model(model_path)
        else:
            # 否则初始化新模型
            self._initialize_models()

    def _load_vulnerability_db(self):
        """加载漏洞数据库"""
        # 这里使用模拟数据，实际应用中应从CVE数据库或漏洞情报平台获取
        try:
            with open('vulnerability_db.json', 'r') as f:
                return json.load(f)
        except FileNotFoundError:
            print("漏洞数据库未找到，使用默认模拟数据")
            return self._generate_sample_vulnerability_db()

    def _generate_sample_vulnerability_db(self):
        """生成样本漏洞数据库"""
        return {
            "CVE-2023-1234": {
                "name": "Apache Log4j 远程代码执行漏洞",
                "severity": "critical",
                "cvss_score": 9.8,
                "description": "Apache Log4j存在远程代码执行漏洞，攻击者可通过构造特殊请求触发漏洞",
                "affected_products": ["Apache Log4j 2.x <= 2.14.1"],
                "attack_vector": "network",
                "exploitability_score": 3.9,
                "impact_score": 5.9,
                "cwe_id": "CWE-77"
            },
            "CVE-2023-5678": {
                "name": "Microsoft Exchange Server 远程代码执行漏洞",
                "severity": "high",
                "cvss_score": 8.8,
                "description": "Microsoft Exchange Server存在远程代码执行漏洞，未经身份验证的攻击者可利用此漏洞在目标系统上执行任意代码",
                "affected_products": ["Exchange Server 2013", "Exchange Server 2016", "Exchange Server 2019"],
                "attack_vector": "network",
                "exploitability_score": 3.9,
                "impact_score": 5.9,
                "cwe_id": "CWE-20"
            },
            "CVE-2023-9101": {
                "name": "Linux内核权限提升漏洞",
                "severity": "high",
                "cvss_score": 7.8,
                "description": "Linux内核存在权限提升漏洞，本地攻击者可利用此漏洞获取root权限",
                "affected_products": ["Linux Kernel 5.4.x <= 5.4.235", "5.10.x <= 5.10.175", "5.15.x <= 5.15.99", "6.1.x <= 6.1.12"],
                "attack_vector": "local",
                "exploitability_score": 1.8,
                "impact_score": 5.9,
                "cwe_id": "CWE-276"
            },
            # 添加更多漏洞...
            "CVE-2023-2467": {
                "name": "Nginx 缓冲区溢出漏洞",
                "severity": "medium",
                "cvss_score": 6.4,
                "description": "Nginx存在缓冲区溢出漏洞，远程攻击者可利用此漏洞造成拒绝服务或执行代码",
                "affected_products": ["Nginx 1.21.0 - 1.21.5"],
                "attack_vector": "network",
                "exploitability_score": 3.4,
                "impact_score": 2.5,
                "cwe_id": "CWE-121"
            },
            "CVE-2023-3890": {
                "name": "MySQL 身份验证绕过漏洞",
                "severity": "critical",
                "cvss_score": 9.3,
                "description": "MySQL存在身份验证绕过漏洞，远程攻击者可利用此漏洞未经授权访问数据库",
                "affected_products": ["MySQL Server 5.7.37", "8.0.28"],
                "attack_vector": "network",
                "exploitability_score": 3.9,
                "impact_score": 5.4,
                "cwe_id": "CWE-287"
            }
        }

    def _initialize_models(self):
        """初始化机器学习模型"""
        # 初始化漏洞检测模型
        self.vulnerability_model = RandomForestClassifier(
            n_estimators=100,
            max_depth=10,
            min_samples_split=5,
            random_state=42
        )

        # 生成样本数据训练模型
        self._train_vulnerability_model()

    def _train_vulnerability_model(self):
        """训练漏洞检测模型"""
        # 生成样本训练数据
        X, y = self._generate_sample_training_data()
        X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

        # 特征缩放
        self.scaler.fit(X_train)
        X_train_scaled = self.scaler.transform(X_train)
        X_test_scaled = self.scaler.transform(X_test)

        # 训练模型
        self.vulnerability_model.fit(X_train_scaled, y_train)

        # 评估模型
        y_pred = self.vulnerability_model.predict(X_test_scaled)
        accuracy = accuracy_score(y_test, y_pred)
        print(f"漏洞检测模型训练完成，准确率: {accuracy:.2f}")

    def _generate_sample_training_data(self):
        """生成样本训练数据"""
        # 特征: [cvss_score, exploitability_score, impact_score, is_network_vector, has_exploit, age_days]
        # 标签: 0-不可利用, 1-可利用
        X = []
        y = []

        for cve_id, vuln_info in self.vulnerability_db.items():
            # 基础特征
            cvss_score = vuln_info['cvss_score']
            exploitability = vuln_info['exploitability_score']
            impact = vuln_info['impact_score']
            is_network = 1 if vuln_info['attack_vector'] == 'network' else 0
            has_exploit = random.choice([0, 1])  # 模拟是否有公开利用代码
            age_days = random.randint(1, 365)  # 模拟漏洞存在时间

            # 添加到特征集
            X.append([cvss_score, exploitability, impact, is_network, has_exploit, age_days])

            # 生成标签: 高CVSS且有利用代码的漏洞更可能被利用
            if cvss_score >= 7.0 and has_exploit == 1:
                y.append(1)  # 可利用
            else:
                # 随机生成一些标签，加入噪声
                y.append(1 if random.random() < 0.3 else 0)

        # 添加更多随机样本
        for _ in range(200):
            cvss_score = random.uniform(0, 10)
            exploitability = random.uniform(0, 4)
            impact = random.uniform(0, 6)
            is_network = random.choice([0, 1])
            has_exploit = random.choice([0, 1])
            age_days = random.randint(1, 365)

            X.append([cvss_score, exploitability, impact, is_network, has_exploit, age_days])

            # 标签生成逻辑
            if cvss_score >= 7.0 and has_exploit == 1:
                y.append(1)
            else:
                y.append(1 if random.random() < 0.15 else 0)

        return np.array(X), np.array(y)

    def _load_model(self, model_path):
        """加载预训练模型"""
        try:
            with open(model_path, 'rb') as f:
                model_data = pickle.load(f)
                self.vulnerability_model = model_data['vulnerability_model']
                self.scaler = model_data['scaler']
            print(f"成功加载预训练模型: {model_path}")
        except Exception as e:
            print(f"模型加载失败: {str(e)}")
            print("使用新初始化的模型")
            self._initialize_models()

    def save_model(self, model_path):
        """保存模型"""
        model_data = {
            'vulnerability_model': self.vulnerability_model,
            'scaler': self.scaler
        }
        with open(model_path, 'wb') as f:
            pickle.dump(model_data, f)
        print(f"模型已保存至: {model_path}")

    def scan_network(self, target_range, scan_type='full'):
        """扫描目标网络"""
        print(f"开始扫描网络: {target_range}, 扫描类型: {scan_type}")
        self.scan_results = {}

        # 模拟网络扫描结果
        # 实际应用中应集成nmap, nessus等扫描工具
        hosts = [f"192.168.1.{i}" for i in range(1, 6)]  # 模拟5个主机

        for host in tqdm(hosts, desc="扫描进度"):
            # 模拟主机扫描结果
            self.scan_results[host] = {
                'status': 'up',
                'os': random.choice(['Linux', 'Windows Server 2019', 'Windows 10', 'macOS']),
                'open_ports': self._generate_open_ports(),
                'services': self._generate_services(host),
                'vulnerabilities': self._detect_vulnerabilities(host)
            }
            time.sleep(0.1)  # 模拟扫描延迟

        print(f"网络扫描完成，发现 {len(hosts)} 个活动主机")
        return self.scan_results

    def _generate_open_ports(self):
        """生成模拟开放端口"""
        common_ports = [21, 22, 23, 25, 80, 8080, 443, 3306, 1433, 5432, 27017]
        open_ports = random.sample(common_ports, k=random.randint(2, 6))
        return {port: 'open' for port in open_ports}

    def _generate_services(self, host):
        """生成模拟服务信息"""
        services = {}
        open_ports = list(self.scan_results[host]['open_ports'].keys()) if host in self.scan_results else []

        for port in open_ports:
            if port == 80 or port == 8080:
                services[port] = {
                    'name': 'http',
                    'product': random.choice(['Apache httpd', 'nginx', 'Microsoft IIS']),
                    'version': self._generate_version(),
                    'cpe': f"cpe:/a:{random.choice(['apache:http_server', 'nginx:nginx', 'microsoft:iis'])}"
                }
            elif port == 443:
                services[port] = {
                    'name': 'https',
                    'product': random.choice(['Apache httpd', 'nginx', 'Microsoft IIS']),
                    'version': self._generate_version(),
                    'ssl_cert': {
                        'subject': {'CN': f"{host}.example.com"},
                        'issuer': 'Let\'s Encrypt'
                    }
                }
            elif port == 22:
                services[port] = {
                    'name': 'ssh',
                    'product': 'OpenSSH',
                    'version': self._generate_version(start=7, end=9),
                    'protocol': 'ssh-2.0'
                }
            elif port == 3306:
                services[port] = {
                    'name': 'mysql',
                    'product': 'MySQL',
                    'version': self._generate_version(start=5, end=8),
                    'protocol': 'tcp'
                }
            elif port == 21:
                services[port] = {
                    'name': 'ftp',
                    'product': random.choice(['vsftpd', 'ProFTPD', 'FileZilla ftpd']),
                    'version': self._generate_version()
                }
            else:
                services[port] = {
                    'name': 'unknown',
                    'product': 'unknown',
                    'version': 'unknown'
                }

        return services

    def _generate_version(self, start=1, end=2):
        """生成模拟版本号"""
        major = random.randint(start, end)
        minor = random.randint(0, 20)
        patch = random.randint(0, 100)
        return f"{major}.{minor}.{patch}"

    def _detect_vulnerabilities(self, host):
        """检测主机漏洞"""
        vulnerabilities = []
        services = self.scan_results[host]['services']

        for port, service in services.items():
            if service['product'] == 'Apache httpd' and port in [80, 443, 8080]:
                # 模拟Apache漏洞检测
                if random.random() < 0.3:
                    vulnerabilities.append({
                        'cve_id': 'CVE-2023-1234',
                        'severity': 'critical',
                        'description': 'Apache Log4j远程代码执行漏洞',
                        'cvss_score': 9.8,
                        'service': f"{service['product']} {service['version']}",
                        'port': port
                    })
            elif service['product'] == 'nginx' and port in [80, 443, 8080]:
                # 模拟Nginx漏洞检测
                if random.random() < 0.2:
                    vulnerabilities.append({
                        'cve_id': 'CVE-2023-2467',
                        'severity': 'medium',
                        'description': 'Nginx缓冲区溢出漏洞',
                        'cvss_score': 6.4,
                        'service': f"{service['product']} {service['version']}",
                        'port': port
                    })
            elif service['product'] == 'MySQL':
                # 模拟MySQL漏洞检测
                if random.random() < 0.25:
                    vulnerabilities.append({
                        'cve_id': 'CVE-2023-3890',
                        'severity': 'critical',
                        'description': 'MySQL身份验证绕过漏洞',
                        'cvss_score': 9.3,
                        'service': f"{service['product']} {service['version']}",
                        'port': port
                    })
            elif service['name'] == 'ssh' and service['product'] == 'OpenSSH':
                # 模拟SSH漏洞检测
                if random.random() < 0.15:
                    vulnerabilities.append({
                        'cve_id': 'CVE-2023-5678',
                        'severity': 'high',
                        'description': 'OpenSSH权限提升漏洞',
                        'cvss_score': 8.8,
                        'service': f"{service['product']} {service['version']}",
                        'port': port
                    })

        return vulnerabilities

    def build_attack_graph(self):
        """构建攻击图"""
        print("开始构建攻击图...")
        self.attack_graph.clear()

        # 添加主机节点
        for host in self.scan_results:
            self.attack_graph.add_node(host, type='host', os=self.scan_results[host]['os'])

        # 添加漏洞节点并连接
        for host in self.scan_results:
            for vuln in self.scan_results[host]['vulnerabilities']:
                vuln_node = f"{host}:{vuln['cve_id']}"
                self.attack_graph.add_node(vuln_node, type='vulnerability', severity=vuln['severity'])
                self.attack_graph.add_edge(host, vuln_node, relationship='has_vulnerability')

                # 基于漏洞类型添加可能的攻击路径
                if vuln['severity'] in ['critical', 'high'] and vuln['cvss_score'] >= 8.0:
                    # 高风险漏洞可以作为攻击跳板
                    for target_host in self.scan_results:
                        if target_host != host:
                            self.attack_graph.add_edge(vuln_node, target_host, relationship='can_attack', probability=self._calculate_attack_probability(vuln))

        print(f"攻击图构建完成，节点数: {len(self.attack_graph.nodes)}, 边数: {len(self.attack_graph.edges)}")
        return self.attack_graph

    def _calculate_attack_probability(self, vuln):
        """计算攻击成功率"""
        # 基于CVSS分数和漏洞类型计算攻击成功率
        base_prob = min(1.0, vuln['cvss_score'] / 10.0)
        if vuln['severity'] == 'critical':
            return min(1.0, base_prob + 0.2)
        elif vuln['severity'] == 'medium':
            return max(0.1, base_prob - 0.2)
        else:
            return max(0.05, base_prob - 0.4)

    def find_attack_paths(self, start_host=None, target_host=None):
        """寻找攻击路径"""
        if not self.attack_graph or len(self.attack_graph.nodes) == 0:
            print("请先构建攻击图")
            self.build_attack_graph()

        print("开始寻找攻击路径...")
        self.attack_paths = []

        # 如果未指定起始和目标主机，则自动选择
        if not start_host:
            start_host = random.choice(list(self.scan_results.keys()))
        if not target_host:
            target_host = random.choice([h for h in self.scan_results if h != start_host])

        print(f"寻找从 {start_host} 到 {target_host} 的攻击路径...")

        # 使用BFS寻找所有可能的路径
        visited = set()
        queue = deque([(start_host, [start_host])])

        while queue:
            current_node, path = queue.popleft()

            if current_node == target_host and len(path) > 1:
                self.attack_paths.append(path)
                continue

            if current_node not in visited:
                visited.add(current_node)
                for neighbor in self.attack_graph.neighbors(current_node):
                    if neighbor not in visited:
                        new_path = path.copy()
                        new_path.append(neighbor)
                        queue.append((neighbor, new_path))

        # 按路径长度排序，优先选择较短路径
        self.attack_paths.sort(key=lambda x: len(x))

        # 过滤重复和不合理的路径
        self.attack_paths = self._filter_attack_paths()

        print(f"找到 {len(self.attack_paths)} 条可能的攻击路径")
        return self.attack_paths

    def _filter_attack_paths(self):
        """过滤攻击路径"""
        filtered = []
        seen = set()

        for path in self.attack_paths:
            # 提取主机序列（排除漏洞节点）
            host_path = [node for node in path if ':' not in node]
            # 转换为元组以便哈希
            host_tuple = tuple(host_path)

            if host_tuple not in seen and len(host_path) >= 2:
                seen.add(host_tuple)
                filtered.append(path)

        # 限制最多返回10条路径
        return filtered[:10]

    def predict_exploitability(self, vulnerability):
        """预测漏洞可利用性"""
        if not self.vulnerability_model:
            print("漏洞检测模型未初始化")
            return 0.0

        # 获取CVE详细信息
        cve_id = vulnerability['cve_id']
        cve_info = self.vulnerability_db.get(cve_id, {})

        # 构建特征向量
        cvss_score = vulnerability.get('cvss_score', cve_info.get('cvss_score', 5.0))
        exploitability = cve_info.get('exploitability_score', 2.0)
        impact = cve_info.get('impact_score', 3.0)
        is_network = 1 if cve_info.get('attack_vector') == 'network' else 0
        has_exploit = 1 if random.random() < 0.7 else 0  # 模拟是否有公开利用
        age_days = random.randint(1, 365)

        features = np.array([[cvss_score, exploitability, impact, is_network, has_exploit, age_days]])
        features_scaled = self.scaler.transform(features)

        # 预测可利用性概率
        probability = self.vulnerability_model.predict_proba(features_scaled)[0][1]

        return round(probability, 2)

    def generate_penetration_test_report(self, output_path='penetration_test_report.md'):
        """生成渗透测试报告"""
        if not self.scan_results or not self.attack_paths:
            print("请先执行扫描和攻击路径分析")
            return

        print(f"生成渗透测试报告: {output_path}")

        # 报告内容
        report = f"# AI辅助渗透测试报告\n\n"
        report += f"## 执行摘要\n\n"
        report += f"- 测试日期: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n"
        report += f"- 目标网络: 模拟网络环境\n"
        report += f"- 活动主机: {len(self.scan_results)}\n"
        report += f"- 发现漏洞: {sum(len(host_info['vulnerabilities']) for host_info in self.scan_results.values())}\n"
        report += f"- 高危漏洞: {sum(1 for host_info in self.scan_results.values() for vuln in host_info['vulnerabilities'] if vuln['severity'] in ['critical', 'high'])}\n"
        report += f"- 发现攻击路径: {len(self.attack_paths)}\n\n"

        report += "## 网络扫描结果\n\n"
        for host, info in self.scan_results.items():
            report += f"### {host}\n"
            report += f"- 操作系统: {info['os']}\n"
            report += f"- 开放端口: {', '.join(map(str, info['open_ports'].keys()))}\n"
            report += "- 服务信息:\n"
            for port, service in info['services'].items():
                report += f"  - 端口 {port}: {service['product']} {service['version']}\n"
            report += "\n"

        report += "## 漏洞摘要\n\n"
        report += "| CVE ID | 主机 | 服务 | 严重级别 | CVSS评分 | 可利用性预测 |\n"
        report += "|--------|------|------|----------|----------|--------------|\n"
        for host, info in self.scan_results.items():
            for vuln in info['vulnerabilities']:
                exploitability = self.predict_exploitability(vuln)
                report += f"| {vuln['cve_id']} | {host} | {vuln['service']} | {vuln['severity']} | {vuln['cvss_score']} | {exploitability} |\n"
        report += "\n"

        report += "## 攻击路径分析\n\n"
        for i, path in enumerate(self.attack_paths, 1):
            report += f"### 路径 #{i}\n"
            report += f"{' → '.join(path)}\n"
            report += f"- 路径长度: {len(path)}\n"
            report += f"- 高风险节点: {sum(1 for node in path if self.attack_graph.nodes.get(node, {}).get('severity') in ['critical', 'high'])}\n\n"

        report += "## 修复建议\n\n"
        report += "1. **高危漏洞优先修复**\n"
        report += "   - 立即修复所有CVSS评分≥9.0的严重漏洞\n"
        report += "   - 对远程代码执行漏洞实施临时缓解措施\n\n"
        report += "2. **网络分段与访问控制**\n"
        report += "   - 实施网络分段，限制不同区域间通信\n"
        report += "   - 对关键服务器实施最小权限原则\n\n"
        report += "3. **安全配置加固**\n"
        report += "   - 关闭不必要的服务和端口\n"
        report += "   - 应用最新安全补丁和更新\n"
        report += "   - 使用强密码策略并实施多因素认证\n\n"
        report += "4. **持续监控与检测**\n"
        report += "   - 部署入侵检测/防御系统(IDS/IPS)\n"
        report += "   - 实施日志集中管理和异常检测\n"
        report += "   - 定期进行漏洞扫描和渗透测试\n\n"

        report += "## 附录\n\n"
        report += "- 测试工具: AI辅助安全渗透测试框架 v1.0\n"
        report += "- 测试方法: 自动化漏洞扫描 + AI攻击路径分析\n"
        report += "- 报告生成日期: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n"

        # 保存报告
        with open(output_path, 'w', encoding='utf-8') as f:
            f.write(report)

        print(f"报告生成成功: {os.path.abspath(output_path)}")
        return output_path

# 主程序演示
if __name__ == '__main__':
    # 创建AI安全测试器实例
    ai_tester = AISecurityTester()

    # 扫描网络
    scan_results = ai_tester.scan_network("192.168.1.0/24")

    # 构建攻击图
    attack_graph = ai_tester.build_attack_graph()

    # 寻找攻击路径
    attack_paths = ai_tester.find_attack_paths()

    # 生成渗透测试报告
    report_path = ai_tester.generate_penetration_test_report()

    # 显示结果摘要
    print("\n=== 测试结果摘要 ===")
    print(f"扫描主机数量: {len(scan_results)}")
    total_vulns = sum(len(host_info['vulnerabilities']) for host_info in scan_results.values())
    print(f"发现漏洞总数: {total_vulns}")
    high_risk_vulns = sum(1 for host_info in scan_results.values() for vuln in host_info['vulnerabilities'] if vuln['severity'] in ['critical', 'high'])
    print(f"高危漏洞数量: {high_risk_vulns}")
    print(f"发现攻击路径数量: {len(attack_paths)}")
    print(f"测试报告已保存至: {report_path}")

    # 可视化攻击图（可选）
    if len(attack_graph.nodes) > 0:
        plt.figure(figsize=(12, 8))
        pos = nx.spring_layout(attack_graph, k=0.3)
        node_colors = []
        for node in attack_graph.nodes:
            if attack_graph.nodes[node]['type'] == 'vulnerability':
                if attack_graph.nodes[node]['severity'] == 'critical':
                    node_colors.append('red')
                elif attack_graph.nodes[node]['severity'] == 'high':
                    node_colors.append('orange')
                else:
                    node_colors.append('yellow')
            else:
                node_colors.append('skyblue')
        nx.draw(attack_graph, pos, with_labels=True, node_color=node_colors, node_size=1500, font_size=8)
        plt.title('AI生成的攻击图可视化')
        plt.savefig('attack_graph.png', dpi=300, bbox_inches='tight')
        print("攻击图可视化已保存至: attack_graph.png")