AI驱动的安全测试：智能漏洞检测与防护

传统安全测试 → 挑战：效率低/覆盖不足/误报高/响应慢/人工依赖 → AI驱动安全测试 → 优势：智能检测/自动分析/动态防护/高效响应/降低依赖

目录
├── 第一章：安全测试的基本概念与挑战
├── 第二章：AI在漏洞智能检测中的应用
├── 第三章：AI在安全风险智能分析中的应用
├── 第四章：AI在防护策略智能生成中的应用
├── 第五章：AI驱动的安全测试框架设计
├── 第六章：实践案例与最佳实践
└── 第七章：未来发展与技能培养

安全测试 = 漏洞检测 + 风险评估 + 安全验证 + 防护建议

传统安全测试挑战分布：漏洞检测效率低(25%) | 误报率高(20%) | 覆盖不足(18%) | 响应速度慢(15%) | 人工依赖(12%) | 复杂系统测试难(10%)

AI技术 → 安全测试变革 → 效率提升 → 准确性提高 → 覆盖增强 → 响应加速 → 成本降低

import os
import re
import numpy as np
import pandas as pd
import tensorflow as tf
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Dense, Dropout, LSTM, Embedding, Bidirectional
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import LabelEncoder
from sklearn.metrics import classification_report, confusion_matrix
import matplotlib.pyplot as plt
from collections import Counter
import nltk
from nltk.tokenize import word_tokenize, sent_tokenize

# 确保中文显示正常
plt.rcParams["font.family"] = ["SimHei", "WenQuanYi Micro Hei", "Heiti TC"]
plt.rcParams["axes.unicode_minus"] = False

# 下载必要的NLTK资源
try:
    nltk.data.find('tokenizers/punkt')
except LookupError:
    nltk.download('punkt')

try:
    nltk.data.find('corpora/stopwords')
except LookupError:
    nltk.download('stopwords')

class CodeVulnerabilityDetector:
    def __init__(self):
        # 初始化参数
        self.max_sequence_length = 200
        self.embedding_dim = 100
        self.vocab_size = 10000
        self.model = None
        self.tokenizer = None
        self.label_encoder = LabelEncoder()
        
        # 创建必要的目录
        self.data_dir = "data"
        self.models_dir = "models"
        self.results_dir = "results"
        
        for dir_path in [self.data_dir, self.models_dir, self.results_dir]:
            if not os.path.exists(dir_path):
                os.makedirs(dir_path)
    
    def load_sample_data(self):
        """加载样本数据"""
        # 创建样本代码数据（实际应用中应该从文件或数据库中加载）
        sample_data = [
            {"code": "def insecure_function(user_input):\n    # 潜在的SQL注入漏洞\n    query = 'SELECT * FROM users WHERE username = " + user_input + "'\n    return execute_query(query)", "label": "SQL Injection"},
            {"code": "def buffer_overflow_example():\n    # 潜在的缓冲区溢出漏洞\n    buffer = [0] * 10\n    for i in range(20):\n        buffer[i] = i\n    return buffer", "label": "Buffer Overflow"},
            {"code": "def xss_vulnerability(user_input):\n    # 潜在的XSS漏洞\n    response = '<div>' + user_input + '</div>'\n    return response", "label": "XSS"},
            {"code": "def safe_function(user_input):\n    # 安全的实现\n    sanitized_input = sanitize(user_input)\n    query = 'SELECT * FROM users WHERE username = %s'\n    return execute_query(query, [sanitized_input])", "label": "Safe"},
            {"code": "def command_injection(user_input):\n    # 潜在的命令注入漏洞\n    os.system('ls -la ' + user_input)\n    return True", "label": "Command Injection"},
            {"code": "def authentication_issue(username, password):\n    # 潜在的认证问题\n    if username == 'admin' and password == 'password':\n        return True\n    return False", "label": "Weak Authentication"},
            {"code": "def encryption_issue(data):\n    # 潜在的加密问题\n    import base64\n    return base64.b64encode(data.encode())", "label": "Insecure Encryption"},
            {"code": "def proper_authentication(username, password):\n    # 安全的认证实现\n    hashed_password = hash_password(password)\n    stored_password = get_stored_password(username)\n    return check_password_hash(hashed_password, stored_password)", "label": "Safe"}
        ]
        
        # 保存样本数据到文件
        df = pd.DataFrame(sample_data)
        data_file_path = os.path.join(self.data_dir, "sample_vulnerability_data.csv")
        df.to_csv(data_file_path, index=False, encoding='utf-8')
        
        print(f"样本数据已保存到: {data_file_path}")
        return df
    
    def preprocess_code(self, code):
        """预处理代码文本"""
        # 移除注释
        code = re.sub(r'#.*$', '', code, flags=re.MULTILINE)
        
        # 移除多余的空白字符
        code = re.sub(r'\s+', ' ', code)
        
        # 规范化代码（转换为小写，移除特殊字符等）
        code = code.lower()
        
        # 分词
        tokens = word_tokenize(code)
        
        return ' '.join(tokens)
    
    def prepare_data(self, df):
        """准备训练数据"""
        # 预处理代码
        df['processed_code'] = df['code'].apply(self.preprocess_code)
        
        # 构建词汇表和Tokenizer
        from tensorflow.keras.preprocessing.text import Tokenizer
        from tensorflow.keras.preprocessing.sequence import pad_sequences
        
        self.tokenizer = Tokenizer(num_words=self.vocab_size, oov_token="<OOV>")
        self.tokenizer.fit_on_texts(df['processed_code'])
        
        # 将文本转换为序列
        sequences = self.tokenizer.texts_to_sequences(df['processed_code'])
        
        # 填充序列到固定长度
        padded_sequences = pad_sequences(sequences, maxlen=self.max_sequence_length, padding='post', truncating='post')
        
        # 编码标签
        labels = self.label_encoder.fit_transform(df['label'])
        
        # 划分训练集和测试集
        X_train, X_test, y_train, y_test = train_test_split(padded_sequences, labels, test_size=0.2, random_state=42)
        
        # 如果类别数量大于2，转换为one-hot编码
        if len(np.unique(labels)) > 2:
            y_train = tf.keras.utils.to_categorical(y_train)
            y_test = tf.keras.utils.to_categorical(y_test)
        
        return X_train, X_test, y_train, y_test
    
    def build_model(self, num_classes):
        """构建深度学习模型"""
        model = Sequential()
        
        # 嵌入层
        model.add(Embedding(input_dim=self.vocab_size, output_dim=self.embedding_dim, input_length=self.max_sequence_length))
        
        # LSTM层
        model.add(Bidirectional(LSTM(64, return_sequences=True)))
        model.add(Dropout(0.2))
        model.add(Bidirectional(LSTM(32)))
        model.add(Dropout(0.2))
        
        # 全连接层
        model.add(Dense(64, activation='relu'))
        model.add(Dropout(0.2))
        
        # 输出层
        if num_classes == 2:
            model.add(Dense(1, activation='sigmoid'))
            loss = 'binary_crossentropy'
        else:
            model.add(Dense(num_classes, activation='softmax'))
            loss = 'categorical_crossentropy'
        
        # 编译模型
        model.compile(loss=loss, optimizer='adam', metrics=['accuracy'])
        
        return model, loss
    
    def train_model(self, X_train, X_test, y_train, y_test):
        """训练模型"""
        # 确定类别数量
        if len(y_train.shape) > 1:
            num_classes = y_train.shape[1]
        else:
            num_classes = len(np.unique(y_train))
        
        # 构建模型
        self.model, loss = self.build_model(num_classes)
        
        # 打印模型结构
        self.model.summary()
        
        # 训练模型
        history = self.model.fit(
            X_train, y_train,
            epochs=10,
            batch_size=32,
            validation_data=(X_test, y_test),
            verbose=1
        )
        
        # 保存模型
        model_path = os.path.join(self.models_dir, "code_vulnerability_model.h5")
        self.model.save(model_path)
        
        print(f"模型已保存到: {model_path}")
        
        # 可视化训练结果
        self.visualize_training(history)
        
        return history
    
    def visualize_training(self, history):
        """可视化训练结果"""
        # 创建图表
        plt.figure(figsize=(12, 4))
        
        # 绘制准确率曲线
        plt.subplot(1, 2, 1)
        plt.plot(history.history['accuracy'], label='训练准确率')
        plt.plot(history.history['val_accuracy'], label='验证准确率')
        plt.title('模型准确率')
        plt.xlabel('Epoch')
        plt.ylabel('准确率')
        plt.legend()
        
        # 绘制损失曲线
        plt.subplot(1, 2, 2)
        plt.plot(history.history['loss'], label='训练损失')
        plt.plot(history.history['val_loss'], label='验证损失')
        plt.title('模型损失')
        plt.xlabel('Epoch')
        plt.ylabel('损失')
        plt.legend()
        
        # 保存图表
        plot_path = os.path.join(self.results_dir, "training_results.png")
        plt.tight_layout()
        plt.savefig(plot_path)
        plt.close()
        
        print(f"训练结果可视化已保存到: {plot_path}")
    
    def evaluate_model(self, X_test, y_test):
        """评估模型"""
        if self.model is None:
            print("错误: 模型未初始化，请先训练模型")
            return
        
        # 评估模型
        loss, accuracy = self.model.evaluate(X_test, y_test, verbose=0)
        print(f"\n模型评估结果:")
        print(f"损失: {loss:.4f}")
        print(f"准确率: {accuracy:.4f}")
        
        # 生成预测结果
        y_pred = self.model.predict(X_test)
        
        # 根据问题类型调整预测结果
        if len(y_test.shape) > 1 and y_test.shape[1] > 1:
            # 多分类问题
            y_pred_classes = np.argmax(y_pred, axis=1)
            y_true_classes = np.argmax(y_test, axis=1)
        else:
            # 二分类问题
            y_pred_classes = (y_pred > 0.5).astype(int).flatten()
            y_true_classes = y_test.flatten()
        
        # 生成分类报告
        print("\n分类报告:")
        report = classification_report(y_true_classes, y_pred_classes, target_names=self.label_encoder.classes_)
        print(report)
        
        # 保存分类报告
        report_path = os.path.join(self.results_dir, "classification_report.txt")
        with open(report_path, 'w', encoding='utf-8') as f:
            f.write(report)
        
        # 生成混淆矩阵
        cm = confusion_matrix(y_true_classes, y_pred_classes)
        self.plot_confusion_matrix(cm)
        
        return loss, accuracy
    
    def plot_confusion_matrix(self, cm):
        """绘制混淆矩阵"""
        plt.figure(figsize=(10, 10))
        
        # 绘制混淆矩阵
        plt.imshow(cm, interpolation='nearest', cmap=plt.cm.Blues)
        plt.title('混淆矩阵')
        plt.colorbar()
        
        # 添加类别标签
        classes = self.label_encoder.classes_
        tick_marks = np.arange(len(classes))
        plt.xticks(tick_marks, classes, rotation=45)
        plt.yticks(tick_marks, classes)
        
        # 添加数值标签
        thresh = cm.max() / 2.
        for i, j in np.ndindex(cm.shape):
            plt.text(j, i, format(cm[i, j], 'd'),
                     horizontalalignment="center",
                     color="white" if cm[i, j] > thresh else "black")
        
        plt.ylabel('真实标签')
        plt.xlabel('预测标签')
        plt.tight_layout()
        
        # 保存混淆矩阵
        cm_path = os.path.join(self.results_dir, "confusion_matrix.png")
        plt.savefig(cm_path)
        plt.close()
        
        print(f"混淆矩阵已保存到: {cm_path}")
    
    def detect_vulnerability(self, code):
        """检测代码中的漏洞"""
        if self.model is None or self.tokenizer is None:
            print("错误: 模型或Tokenizer未初始化，请先训练模型")
            return None
        
        # 预处理代码
        processed_code = self.preprocess_code(code)
        
        # 转换为序列
        sequence = self.tokenizer.texts_to_sequences([processed_code])
        
        # 填充序列
        padded_sequence = pad_sequences(sequence, maxlen=self.max_sequence_length, padding='post', truncating='post')
        
        # 预测
        prediction = self.model.predict(padded_sequence)
        
        # 根据问题类型解析预测结果
        if prediction.shape[1] > 1:
            # 多分类问题
            predicted_class = np.argmax(prediction, axis=1)[0]
        else:
            # 二分类问题
            predicted_class = 1 if prediction[0][0] > 0.5 else 0
        
        # 获取预测的漏洞类型
        vulnerability_type = self.label_encoder.inverse_transform([predicted_class])[0]
        
        # 获取置信度
        if prediction.shape[1] > 1:
            confidence = prediction[0][predicted_class]
        else:
            confidence = prediction[0][0] if predicted_class == 1 else 1 - prediction[0][0]
        
        return {
            "vulnerability_type": vulnerability_type,
            "confidence": float(confidence)
        }

# 主函数
def main():
    print("===== AI驱动的代码漏洞检测工具 =====")
    
    # 创建漏洞检测器实例
    detector = CodeVulnerabilityDetector()
    
    # 加载样本数据
    print("加载样本数据...")
    df = detector.load_sample_data()
    
    # 准备训练数据
    print("准备训练数据...")
    X_train, X_test, y_train, y_test = detector.prepare_data(df)
    
    # 训练模型
    print("训练模型...")
    history = detector.train_model(X_train, X_test, y_train, y_test)
    
    # 评估模型
    print("评估模型...")
    loss, accuracy = detector.evaluate_model(X_test, y_test)
    
    # 测试模型
    print("\n===== 测试代码漏洞检测 =====")
    
    # 测试安全代码
    safe_code = """def safe_function(user_input):
    # 安全的实现
    sanitized_input = sanitize(user_input)
    query = 'SELECT * FROM users WHERE username = %s'
    return execute_query(query, [sanitized_input])
"""
    
    print("\n测试安全代码:")
    print(safe_code)
    result = detector.detect_vulnerability(safe_code)
    print(f"检测结果: {result['vulnerability_type']} (置信度: {result['confidence']:.4f})")
    
    # 测试有漏洞的代码
    vulnerable_code = """def vulnerable_function(user_input):
    # 潜在的SQL注入漏洞
    query = 'SELECT * FROM users WHERE username = " + user_input + "'
    return execute_query(query)
"""
    
    print("\n测试有漏洞的代码:")
    print(vulnerable_code)
    result = detector.detect_vulnerability(vulnerable_code)
    print(f"检测结果: {result['vulnerability_type']} (置信度: {result['confidence']:.4f})")
    
    print("\n工具使用说明：")
    print("1. 在 data 目录中修改 sample_vulnerability_data.csv 文件，添加更多的训练数据")
    print("2. 重新运行本工具，训练更准确的模型")
    print("3. 使用 detector.detect_vulnerability() 函数检测新的代码")
    print("4. 在 results 目录中查看评估结果和可视化图表")

if __name__ == "__main__":
    main()

初始策略 → 执行渗透测试 → 收集反馈 → 分析结果 → 优化策略 → 再次执行 → 持续优化

import numpy as np
import pandas as pd
import networkx as nx
import matplotlib.pyplot as plt
import torch
import torch.nn.functional as F
from torch_geometric.data import Data, DataLoader
from torch_geometric.nn import GCNConv, GATConv
import os
from sklearn.preprocessing import LabelEncoder
from sklearn.model_selection import train_test_split
from sklearn.metrics import classification_report

# 确保中文显示正常
plt.rcParams["font.family"] = ["SimHei", "WenQuanYi Micro Hei", "Heiti TC"]
plt.rcParams["axes.unicode_minus"] = False

class AttackPathAnalyzer:
    def __init__(self):
        # 初始化参数
        self.embedding_dim = 64
        self.hidden_dim = 128
        self.output_dim = 2
        self.learning_rate = 0.01
        self.epochs = 20
        self.model = None
        
        # 创建必要的目录
        self.data_dir = "data"
        self.models_dir = "models"
        self.results_dir = "results"
        
        for dir_path in [self.data_dir, self.models_dir, self.results_dir]:
            if not os.path.exists(dir_path):
                os.makedirs(dir_path)
    
    def generate_sample_attack_graph(self):
        """生成样本攻击图数据"""
        # 创建一个有向图表示系统组件和潜在的攻击路径
        G = nx.DiGraph()
        
        # 添加节点（系统组件）
        components = [
            {"id": 0, "name": "外部网络", "type": "network", "value": 0.1},
            {"id": 1, "name": "防火墙", "type": "security_device", "value": 0.8},
            {"id": 2, "name": "Web服务器", "type": "server", "value": 0.6},
            {"id": 3, "name": "应用服务器", "type": "server", "value": 0.7},
            {"id": 4, "name": "数据库服务器", "type": "server", "value": 0.9},
            {"id": 5, "name": "内部网络", "type": "network", "value": 0.2},
            {"id": 6, "name": "管理员账户", "type": "user", "value": 0.8},
            {"id": 7, "name": "普通用户账户", "type": "user", "value": 0.3}
        ]
        
        for component in components:
            G.add_node(component["id"], **component)
        
        # 添加边（潜在的攻击路径）
        attack_paths = [
            (0, 1, {"name": "绕过防火墙", "probability": 0.3, "exploit_type": "firewall_bypass"}),
            (1, 2, {"name": "Web漏洞攻击", "probability": 0.5, "exploit_type": "web_vulnerability"}),
            (2, 3, {"name": "服务间攻击", "probability": 0.4, "exploit_type": "service_to_service"}),
            (3, 4, {"name": "数据库访问", "probability": 0.6, "exploit_type": "database_access"}),
            (0, 5, {"name": "内网渗透", "probability": 0.2, "exploit_type": "internal_penetration"}),
            (5, 4, {"name": "内网访问数据库", "probability": 0.7, "exploit_type": "internal_db_access"}),
            (6, 4, {"name": "管理员访问", "probability": 0.9, "exploit_type": "admin_access"}),
            (7, 2, {"name": "用户访问Web", "probability": 0.8, "exploit_type": "user_web_access"}),
            (7, 3, {"name": "用户访问应用", "probability": 0.5, "exploit_type": "user_app_access"}),
            (2, 0, {"name": "反向连接", "probability": 0.4, "exploit_type": "reverse_connection"})
        ]
        
        for u, v, attr in attack_paths:
            G.add_edge(u, v, **attr)
        
        # 保存图数据到文件
        graph_data = {
            "nodes": components,
            "edges": attack_paths
        }
        
        # 保存节点数据
        nodes_df = pd.DataFrame(components)
        nodes_file_path = os.path.join(self.data_dir, "nodes.csv")
        nodes_df.to_csv(nodes_file_path, index=False, encoding='utf-8')
        
        # 保存边数据
        edges_data = []
        for u, v, attr in attack_paths:
            edge_data = {"source": u, "target": v}
            edge_data.update(attr)
            edges_data.append(edge_data)
        
        edges_df = pd.DataFrame(edges_data)
        edges_file_path = os.path.join(self.data_dir, "edges.csv")
        edges_df.to_csv(edges_file_path, index=False, encoding='utf-8')
        
        print(f"样本攻击图数据已保存到: {nodes_file_path} 和 {edges_file_path}")
        
        # 可视化攻击图
        self.visualize_attack_graph(G)
        
        return G, nodes_df, edges_df
    
    def visualize_attack_graph(self, G):
        """可视化攻击图"""
        plt.figure(figsize=(12, 10))
        
        # 设置节点位置
        pos = nx.spring_layout(G, k=0.3, iterations=50)
        
        # 获取节点类型和值用于可视化
        node_types = nx.get_node_attributes(G, 'type')
        node_values = nx.get_node_attributes(G, 'value')
        node_names = nx.get_node_attributes(G, 'name')
        
        # 设置节点颜色和大小
        type_colors = {
            'network': 'lightblue',
            'security_device': 'green',
            'server': 'red',
            'user': 'yellow'
        }
        
        node_colors = [type_colors.get(node_types[node], 'gray') for node in G.nodes()]
        node_sizes = [3000 * (node_values.get(node, 0.5) + 0.5) for node in G.nodes()]
        
        # 绘制节点
        nx.draw_networkx_nodes(G, pos, node_color=node_colors, node_size=node_sizes, alpha=0.8)
        
        # 绘制边和标签
        edge_labels = {(u, v): f"{d['name']} ({d['probability']:.1f})" for u, v, d in G.edges(data=True)}
        nx.draw_networkx_edges(G, pos, edge_color='gray', arrows=True, arrowstyle='->', alpha=0.6)
        nx.draw_networkx_labels(G, pos, labels=node_names, font_size=10, font_weight='bold')
        nx.draw_networkx_edge_labels(G, pos, edge_labels=edge_labels, font_size=8, alpha=0.7)
        
        # 添加图例
        from matplotlib.lines import Line2D
        legend_elements = [
            Line2D([0], [0], marker='o', color='w', label='网络', markerfacecolor='lightblue', markersize=10),
            Line2D([0], [0], marker='o', color='w', label='安全设备', markerfacecolor='green', markersize=10),
            Line2D([0], [0], marker='o', color='w', label='服务器', markerfacecolor='red', markersize=10),
            Line2D([0], [0], marker='o', color='w', label='用户', markerfacecolor='yellow', markersize=10)
        ]
        plt.legend(handles=legend_elements, loc='best')
        
        plt.title('系统攻击路径图')
        plt.axis('off')
        
        # 保存可视化结果
        viz_path = os.path.join(self.results_dir, "attack_graph.png")
        plt.tight_layout()
        plt.savefig(viz_path)
        plt.close()
        
        print(f"攻击图可视化已保存到: {viz_path}")
    
    def prepare_graph_data(self, G, nodes_df, edges_df):
        """准备图神经网络的输入数据"""
        # 提取节点特征
        node_features = []
        for _, node in nodes_df.iterrows():
            # 这里简化处理，实际应用中应该使用更丰富的特征
            # 节点类型的独热编码
            type_encoding = [0] * 4  # 假设有4种节点类型
            type_mapping = {'network': 0, 'security_device': 1, 'server': 2, 'user': 3}
            type_idx = type_mapping.get(node['type'], 0)
            type_encoding[type_idx] = 1
            
            # 节点值（表示重要性或脆弱性）
            value = [node['value']]
            
            # 组合特征
            feature = type_encoding + value
            node_features.append(feature)
        
        node_features = torch.tensor(node_features, dtype=torch.float)
        
        # 提取边索引和边特征
        edge_index = []
        edge_features = []
        
        for _, edge in edges_df.iterrows():
            edge_index.append([edge['source'], edge['target']])
            
            # 边特征：攻击概率和攻击类型
            # 简化处理，使用攻击概率作为主要特征
            edge_feature = [edge['probability']]
            edge_features.append(edge_feature)
        
        edge_index = torch.tensor(edge_index, dtype=torch.long).t().contiguous()
        edge_features = torch.tensor(edge_features, dtype=torch.float)
        
        # 创建标签：假设我们要预测哪些边是高风险的攻击路径（概率>0.5）
        labels = []
        for _, edge in edges_df.iterrows():
            labels.append(1 if edge['probability'] > 0.5 else 0)
        
        labels = torch.tensor(labels, dtype=torch.long)
        
        # 创建图数据对象
        data = Data(
            x=node_features,
            edge_index=edge_index,
            edge_attr=edge_features,
            y=labels
        )
        
        return data
    
    def build_gcn_model(self, input_dim, edge_attr_dim):
        """构建图卷积网络模型"""
        class GCNModel(torch.nn.Module):
            def __init__(self, input_dim, edge_attr_dim, embedding_dim, hidden_dim, output_dim):
                super(GCNModel, self).__init__()
                # 节点嵌入层
                self.node_embedding = torch.nn.Linear(input_dim, embedding_dim)
                # 边嵌入层
                self.edge_embedding = torch.nn.Linear(edge_attr_dim, embedding_dim)
                # GCN层
                self.conv1 = GCNConv(embedding_dim, hidden_dim)
                self.conv2 = GCNConv(hidden_dim, hidden_dim)
                # 输出层
                self.out = torch.nn.Linear(hidden_dim * 2, output_dim)  # 乘以2因为我们需要连接源节点和目标节点的特征
            
            def forward(self, data):
                x, edge_index, edge_attr, edge_index_2 = data.x, data.edge_index, data.edge_attr, data.edge_index
                
                # 嵌入节点和边特征
                x = self.node_embedding(x)
                edge_attr_emb = self.edge_embedding(edge_attr)
                
                # 应用GCN层
                x = self.conv1(x, edge_index)
                x = F.relu(x)
                x = F.dropout(x, training=self.training)
                x = self.conv2(x, edge_index)
                x = F.relu(x)
                
                # 对于每条边，获取源节点和目标节点的特征并连接
                # 注意：这里我们假设edge_index的形状是[2, num_edges]，其中第一行是源节点索引，第二行是目标节点索引
                src, dst = edge_index_2
                edge_rep = torch.cat([x[src], x[dst]], dim=1)
                
                # 应用输出层
                out = self.out(edge_rep)
                
                return out
        
        # 创建模型实例
        model = GCNModel(input_dim, edge_attr_dim, self.embedding_dim, self.hidden_dim, self.output_dim)
        return model
    
    def train_model(self, data):
        """训练图神经网络模型"""
        # 准备数据加载器
        loader = DataLoader([data], batch_size=1)
        
        # 确定输入维度
        input_dim = data.x.shape[1]
        edge_attr_dim = data.edge_attr.shape[1] if data.edge_attr is not None else 1
        
        # 构建模型
        self.model = self.build_gcn_model(input_dim, edge_attr_dim)
        
        # 定义优化器和损失函数
        optimizer = torch.optim.Adam(self.model.parameters(), lr=self.learning_rate)
        criterion = torch.nn.CrossEntropyLoss()
        
        # 训练模型
        self.model.train()
        for epoch in range(self.epochs):
            total_loss = 0
            for batch in loader:
                optimizer.zero_grad()
                out = self.model(batch)
                loss = criterion(out, batch.y)
                loss.backward()
                optimizer.step()
                total_loss += loss.item()
            
            if (epoch + 1) % 5 == 0:
                print(f'Epoch {epoch+1}/{self.epochs}, Loss: {total_loss:.4f}')
        
        # 保存模型
        model_path = os.path.join(self.models_dir, "attack_path_model.pth")
        torch.save(self.model.state_dict(), model_path)
        
        print(f"模型已保存到: {model_path}")
        
        return self.model
    
    def evaluate_model(self, data):
        """评估模型"""
        if self.model is None:
            print("错误: 模型未初始化，请先训练模型")
            return
        
        # 评估模型
        self.model.eval()
        with torch.no_grad():
            out = self.model(data)
            pred = out.argmax(dim=1)
            
            # 生成分类报告
            print("\n分类报告:")
            report = classification_report(data.y.numpy(), pred.numpy(), target_names=["低风险", "高风险"])
            print(report)
            
            # 保存分类报告
            report_path = os.path.join(self.results_dir, "classification_report.txt")
            with open(report_path, 'w', encoding='utf-8') as f:
                f.write(report)
            
            # 可视化预测结果
            self.visualize_prediction_results(data, pred)
        
        return pred
    
    def visualize_prediction_results(self, data, pred):
        """可视化预测结果"""
        # 重建图结构
        G = nx.DiGraph()
        
        # 添加节点
        for i in range(data.x.shape[0]):
            G.add_node(i)
        
        # 添加边并标记预测结果
        edge_index = data.edge_index.numpy()
        for i in range(edge_index.shape[1]):
            u = edge_index[0, i]
            v = edge_index[1, i]
            # 真实标签和预测标签
            true_label = data.y[i].item()
            pred_label = pred[i].item()
            
            # 设置边的颜色和宽度
            if true_label == 1 and pred_label == 1:
                # 真阳性
                color = 'red'
                width = 2.0
            elif true_label == 0 and pred_label == 0:
                # 真阴性
                color = 'green'
                width = 1.0
            elif true_label == 1 and pred_label == 0:
                # 假阴性
                color = 'orange'
                width = 1.5
            else:
                # 假阳性
                color = 'blue'
                width = 1.5
            
            G.add_edge(u, v, true_label=true_label, pred_label=pred_label, color=color, width=width)
        
        # 可视化结果
        plt.figure(figsize=(12, 10))
        
        # 设置节点位置
        pos = nx.spring_layout(G, k=0.3, iterations=50)
        
        # 绘制节点
        nx.draw_networkx_nodes(G, pos, node_color='lightblue', node_size=2000, alpha=0.8)
        
        # 绘制边
        edges = G.edges(data=True)
        colors = [edge[2]['color'] for edge in edges]
        widths = [edge[2]['width'] for edge in edges]
        nx.draw_networkx_edges(G, pos, edge_color=colors, width=widths, arrows=True, arrowstyle='->', alpha=0.8)
        
        # 添加节点标签
        nx.draw_networkx_labels(G, pos, font_size=12, font_weight='bold')
        
        # 添加图例
        from matplotlib.lines import Line2D
        legend_elements = [
            Line2D([0], [0], color='red', lw=2, label='真阳性 (高风险且预测正确)'),
            Line2D([0], [0], color='green', lw=2, label='真阴性 (低风险且预测正确)'),
            Line2D([0], [0], color='orange', lw=2, label='假阴性 (高风险但预测为低风险)'),
            Line2D([0], [0], color='blue', lw=2, label='假阳性 (低风险但预测为高风险)')
        ]
        plt.legend(handles=legend_elements, loc='best')
        
        plt.title('攻击路径风险预测结果')
        plt.axis('off')
        
        # 保存可视化结果
        viz_path = os.path.join(self.results_dir, "prediction_results.png")
        plt.tight_layout()
        plt.savefig(viz_path)
        plt.close()
        
        print(f"预测结果可视化已保存到: {viz_path}")
    
    def predict_attack_risk(self, new_data):
        """预测新数据的攻击风险"""
        if self.model is None:
            print("错误: 模型未初始化，请先训练模型")
            return None
        
        # 预测风险
        self.model.eval()
        with torch.no_grad():
            out = self.model(new_data)
            pred = out.argmax(dim=1)
            probabilities = F.softmax(out, dim=1)
        
        return {
            "predictions": pred.numpy(),
            "probabilities": probabilities.numpy()
        }

# 主函数
def main():
    print("===== AI驱动的攻击路径分析工具 =====")
    
    # 创建攻击路径分析器实例
    analyzer = AttackPathAnalyzer()
    
    # 生成样本攻击图数据
    print("生成样本攻击图数据...")
    G, nodes_df, edges_df = analyzer.generate_sample_attack_graph()
    
    # 准备图神经网络数据
    print("准备图神经网络数据...")
    data = analyzer.prepare_graph_data(G, nodes_df, edges_df)
    
    # 训练模型
    print("训练模型...")
    model = analyzer.train_model(data)
    
    # 评估模型
    print("评估模型...")
    pred = analyzer.evaluate_model(data)
    
    # 分析预测结果
    high_risk_edges = torch.sum(pred).item()
    total_edges = len(pred)
    print(f"\n预测结果分析:")
    print(f"高风险攻击路径数量: {high_risk_edges}/{total_edges}")
    
    # 找出高风险路径
    edge_index = data.edge_index.numpy()
    high_risk_paths = []
    for i in range(total_edges):
        if pred[i] == 1:
            u = edge_index[0, i]
            v = edge_index[1, i]
            # 获取节点名称
            u_name = nodes_df.loc[u, 'name']
            v_name = nodes_df.loc[v, 'name']
            high_risk_paths.append((u_name, v_name))
    
    print("\n高风险攻击路径:")
    for path in high_risk_paths:
        print(f"  {path[0]} -> {path[1]}")
    
    print("\n工具使用说明：")
    print("1. 在 data 目录中修改 nodes.csv 和 edges.csv 文件，添加您的系统组件和攻击路径数据")
    print("2. 重新运行本工具，训练更准确的模型")
    print("3. 使用 analyzer.predict_attack_risk() 函数预测新的攻击路径风险")
    print("4. 在 results 目录中查看分析结果和可视化图表")

if __name__ == "__main__":
    main()

数据收集 → 实体识别 → 关系抽取 → 知识图谱构建 → 风险关联分析 → 风险可视化 → 风险预警

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import gym
from gym import spaces
from stable_baselines3 import PPO
from stable_baselines3.common.env_util import make_vec_env
from stable_baselines3.common.vec_env import DummyVecEnv
import os
import random

# 确保中文显示正常
plt.rcParams["font.family"] = ["SimHei", "WenQuanYi Micro Hei", "Heiti TC"]
plt.rcParams["axes.unicode_minus"] = False

class SecurityEnv(gym.Env):
    """安全防护策略优化的强化学习环境"""
    def __init__(self):
        super(SecurityEnv, self).__init__()
        
        # 定义环境参数
        self.num_assets = 5  # 资产数量
        self.num_attacks = 3  # 攻击类型数量
        self.num_defenses = 4  # 防御策略数量
        
        # 初始化资产状态（价值和脆弱性）
        self.asset_values = np.array([10, 20, 15, 25, 30])  # 资产价值
        self.asset_vulnerabilities = np.array([0.8, 0.6, 0.7, 0.5, 0.4])  # 资产脆弱性（0-1，越高越脆弱）
        
        # 定义攻击类型及其特性
        self.attacks = {
            0: {"name": "SQL注入", "success_rate": 0.7, "impact": 0.8, "target_preference": [0.3, 0.1, 0.2, 0.3, 0.1]},
            1: {"name": "XSS攻击", "success_rate": 0.8, "impact": 0.6, "target_preference": [0.2, 0.3, 0.1, 0.2, 0.2]},
            2: {"name": "DDoS攻击", "success_rate": 0.9, "impact": 0.9, "target_preference": [0.1, 0.2, 0.2, 0.2, 0.3]}
        }
        
        # 定义防御策略及其特性
        self.defenses = {
            0: {"name": "Web应用防火墙", "cost": 5, "effectiveness": [0.9, 0.8, 0.1], "coverage": [1, 1, 0, 0, 0]},
            1: {"name": "入侵检测系统", "cost": 4, "effectiveness": [0.7, 0.7, 0.8], "coverage": [1, 1, 1, 1, 1]},
            2: {"name": "定期漏洞扫描", "cost": 3, "effectiveness": [0.8, 0.9, 0.2], "coverage": [1, 1, 1, 1, 1]},
            3: {"name": "DDoS防护", "cost": 6, "effectiveness": [0.1, 0.1, 0.9], "coverage": [0, 1, 1, 1, 1]}
        }
        
        # 定义状态空间（资产状态、攻击历史、防御部署情况）
        # 状态空间：[资产价值 * 资产脆弱性 * 攻击风险] + [防御部署情况] + [最近攻击历史]
        self.observation_space = spaces.Box(
            low=0, high=1, 
            shape=(self.num_assets + self.num_defenses + self.num_attacks,), 
            dtype=np.float32
        )
        
        # 定义动作空间（选择部署或调整的防御策略）
        # 动作空间：[防御策略0是否启用, 防御策略1是否启用, ...]（0表示不启用，1表示启用）
        self.action_space = spaces.MultiBinary(self.num_defenses)
        
        # 初始化状态
        self.state = None
        self.defense_deployment = np.zeros(self.num_defenses, dtype=np.int32)
        self.attack_history = np.zeros(self.num_attacks, dtype=np.float32)
        self.current_step = 0
        self.max_steps = 100
        
        # 创建必要的目录
        self.results_dir = "results"
        if not os.path.exists(self.results_dir):
            os.makedirs(self.results_dir)
    
    def _get_observation(self):
        """获取当前环境的观察状态"""
        # 计算资产风险
        asset_risks = self.asset_vulnerabilities * np.sum([
            self.attacks[i]["success_rate"] * self.attacks[i]["target_preference"] for i in range(self.num_attacks)
        ], axis=0) / self.num_attacks
        asset_risks = asset_risks / np.max(asset_risks) if np.max(asset_risks) > 0 else asset_risks
        
        # 组合观察状态
        observation = np.concatenate([
            asset_risks,
            self.defense_deployment.astype(np.float32),
            self.attack_history
        ])
        
        return observation
    
    def reset(self):
        """重置环境状态"""
        # 重置防御部署
        self.defense_deployment = np.zeros(self.num_defenses, dtype=np.int32)
        
        # 重置攻击历史
        self.attack_history = np.zeros(self.num_attacks, dtype=np.float32)
        
        # 重置步骤计数器
        self.current_step = 0
        
        # 获取初始观察状态
        self.state = self._get_observation()
        
        return self.state
    
    def _simulate_attack(self):
        """模拟攻击发生"""
        # 随机选择攻击类型
        attack_type = random.choice(range(self.num_attacks))
        attack = self.attacks[attack_type]
        
        # 根据目标偏好选择攻击目标
        target_probs = attack["target_preference"]
        target_probs = target_probs / np.sum(target_probs)
        target_asset = np.random.choice(range(self.num_assets), p=target_probs)
        
        # 计算攻击成功率（考虑防御策略的影响）
        success_rate = attack["success_rate"]
        for defense_id in range(self.num_defenses):
            if self.defense_deployment[defense_id] == 1 and self.defenses[defense_id]["coverage"][target_asset] == 1:
                # 应用防御策略的效果
                effectiveness = self.defenses[defense_id]["effectiveness"][attack_type]
                success_rate *= (1 - effectiveness)
        
        # 模拟攻击结果
        attack_successful = random.random() < success_rate
        
        # 计算攻击影响
        if attack_successful:
            impact = attack["impact"] * self.asset_values[target_asset] * self.asset_vulnerabilities[target_asset]
            # 更新攻击历史
            self.attack_history[attack_type] += 1
            self.attack_history = np.clip(self.attack_history / 10, 0, 1)  # 归一化攻击历史
        else:
            impact = 0
        
        return attack_type, target_asset, attack_successful, impact
    
    def _calculate_reward(self, impact):
        """计算奖励值"""
        # 计算防御成本
        defense_cost = np.sum([self.defenses[i]["cost"] * self.defense_deployment[i] for i in range(self.num_defenses)])
        
        # 计算攻击损失
        attack_loss = impact
        
        # 计算总奖励（负的总成本和损失）
        total_cost = defense_cost + attack_loss
        reward = -total_cost
        
        # 额外奖励：成功防御攻击
        if impact == 0 and self.current_step > 0:
            reward += 5  # 成功防御的奖励
        
        # 额外奖励：资源使用效率（部署的防御策略数量适中）
        num_deployed = np.sum(self.defense_deployment)
        if num_deployed > 0 and num_deployed < self.num_defenses:
            reward += 2  # 资源使用效率奖励
        
        return reward
    
    def step(self, action):
        """执行一步动作并返回环境反馈"""
        # 更新防御部署
        self.defense_deployment = action.astype(np.int32)
        
        # 模拟攻击
        attack_type, target_asset, attack_successful, impact = self._simulate_attack()
        
        # 计算奖励
        reward = self._calculate_reward(impact)
        
        # 更新步骤计数器
        self.current_step += 1
        
        # 检查是否达到终止条件
        done = self.current_step >= self.max_steps
        
        # 更新状态
        self.state = self._get_observation()
        
        # 提供额外信息
        info = {
            "attack_type": attack_type,
            "target_asset": target_asset,
            "attack_successful": attack_successful,
            "impact": impact,
            "defense_cost": np.sum([self.defenses[i]["cost"] * self.defense_deployment[i] for i in range(self.num_defenses)]),
            "defense_deployment": self.defense_deployment
        }
        
        return self.state, reward, done, info
    
    def render(self, mode='human'):
        """渲染环境状态"""
        if mode == 'human':
            print(f"\n步骤: {self.current_step}/{self.max_steps}")
            print(f"资产风险: {self.state[:self.num_assets]}")
            print(f"防御部署: {self.defense_deployment}")
            print(f"攻击历史: {self.attack_history}")
    
    def close(self):
        """关闭环境"""
        pass

class SecurityDefenseOptimizer:
    def __init__(self):
        # 创建环境
        self.env = SecurityEnv()
        
        # 创建必要的目录
        self.models_dir = "models"
        self.results_dir = "results"
        
        for dir_path in [self.models_dir, self.results_dir]:
            if not os.path.exists(dir_path):
                os.makedirs(dir_path)
        
        # 初始化模型
        self.model = None
    
    def train_model(self, total_timesteps=10000):
        """训练强化学习模型"""
        # 创建模型
        self.model = PPO("MlpPolicy", self.env, verbose=1)
        
        # 训练模型
        print("开始训练模型...")
        self.model.learn(total_timesteps=total_timesteps)
        
        # 保存模型
        model_path = os.path.join(self.models_dir, "security_defense_model.zip")
        self.model.save(model_path)
        
        print(f"模型已保存到: {model_path}")
        
        return self.model
    
    def evaluate_model(self, num_episodes=10):
        """评估模型性能"""
        if self.model is None:
            print("错误: 模型未初始化，请先训练模型")
            return
        
        # 评估模型
        print(f"\n开始评估模型（{num_episodes}个回合）...")
        
        total_rewards = []
        total_impacts = []
        total_costs = []
        deployment_histories = []
        
        for episode in range(num_episodes):
            obs = self.env.reset()
            episode_reward = 0
            episode_impact = 0
            episode_cost = 0
            deployment_history = []
            
            done = False
            while not done:
                # 使用模型选择动作
                action, _ = self.model.predict(obs, deterministic=True)
                
                # 记录部署历史
                deployment_history.append(action.copy())
                
                # 执行动作
                obs, reward, done, info = self.env.step(action)
                
                # 累积奖励和影响
                episode_reward += reward
                episode_impact += info["impact"]
                episode_cost += info["defense_cost"]
            
            # 记录回合结果
            total_rewards.append(episode_reward)
            total_impacts.append(episode_impact)
            total_costs.append(episode_cost)
            deployment_histories.append(np.array(deployment_history))
            
            print(f"回合 {episode+1}/{num_episodes}: 奖励 = {episode_reward:.2f}, 损失 = {episode_impact:.2f}, 成本 = {episode_cost:.2f}")
        
        # 计算平均结果
        avg_reward = np.mean(total_rewards)
        avg_impact = np.mean(total_impacts)
        avg_cost = np.mean(total_costs)
        
        print(f"\n评估结果汇总:")
        print(f"平均奖励: {avg_reward:.2f}")
        print(f"平均损失: {avg_impact:.2f}")
        print(f"平均成本: {avg_cost:.2f}")
        
        # 保存评估结果
        results = {
            "avg_reward": avg_reward,
            "avg_impact": avg_impact,
            "avg_cost": avg_cost,
            "total_rewards": total_rewards,
            "total_impacts": total_impacts,
            "total_costs": total_costs
        }
        
        results_path = os.path.join(self.results_dir, "evaluation_results.json")
        import json
        with open(results_path, 'w', encoding='utf-8') as f:
            json.dump(results, f, ensure_ascii=False, indent=2)
        
        # 可视化评估结果
        self.visualize_evaluation(results, deployment_histories)
        
        return results
    
    def visualize_evaluation(self, results, deployment_histories):
        """可视化评估结果"""
        # 绘制奖励、损失和成本曲线
        plt.figure(figsize=(15, 5))
        
        # 奖励曲线
        plt.subplot(1, 3, 1)
        plt.plot(results["total_rewards"], marker='o')
        plt.axhline(y=results["avg_reward"], color='r', linestyle='--', label=f'平均奖励: {results["avg_reward"]:.2f}')
        plt.title('每个回合的奖励')
        plt.xlabel('回合')
        plt.ylabel('奖励')
        plt.legend()
        
        # 损失曲线
        plt.subplot(1, 3, 2)
        plt.plot(results["total_impacts"], marker='o')
        plt.axhline(y=results["avg_impact"], color='r', linestyle='--', label=f'平均损失: {results["avg_impact"]:.2f}')
        plt.title('每个回合的损失')
        plt.xlabel('回合')
        plt.ylabel('损失')
        plt.legend()
        
        # 成本曲线
        plt.subplot(1, 3, 3)
        plt.plot(results["total_costs"], marker='o')
        plt.axhline(y=results["avg_cost"], color='r', linestyle='--', label=f'平均成本: {results["avg_cost"]:.2f}')
        plt.title('每个回合的防御成本')
        plt.xlabel('回合')
        plt.ylabel('成本')
        plt.legend()
        
        plt.tight_layout()
        
        # 保存图表
        plot_path = os.path.join(self.results_dir, "evaluation_results.png")
        plt.savefig(plot_path)
        plt.close()
        
        print(f"评估结果可视化已保存到: {plot_path}")
        
        # 可视化防御部署策略
        self.visualize_defense_deployment(deployment_histories)
    
    def visualize_defense_deployment(self, deployment_histories):
        """可视化防御部署策略"""
        # 计算每种防御策略的平均部署率
        avg_deployment = np.mean(np.concatenate(deployment_histories), axis=0)
        
        # 获取防御策略名称
        defense_names = [self.env.defenses[i]["name"] for i in range(self.env.num_defenses)]
        
        # 绘制防御策略部署率条形图
        plt.figure(figsize=(12, 6))
        plt.bar(defense_names, avg_deployment)
        plt.title('防御策略平均部署率')
        plt.xlabel('防御策略')
        plt.ylabel('部署率')
        plt.ylim(0, 1)
        
        # 添加数值标签
        for i, v in enumerate(avg_deployment):
            plt.text(i, v + 0.02, f'{v:.2f}', ha='center')
        
        plt.xticks(rotation=45, ha='right')
        plt.tight_layout()
        
        # 保存图表
        plot_path = os.path.join(self.results_dir, "defense_deployment.png")
        plt.savefig(plot_path)
        plt.close()
        
        print(f"防御部署可视化已保存到: {plot_path}")
    
    def generate_security_policy(self):
        """生成优化的安全策略"""
        if self.model is None:
            print("错误: 模型未初始化，请先训练模型")
            return None
        
        # 模拟多个场景，找出最佳策略
        obs = self.env.reset()
        action, _ = self.model.predict(obs, deterministic=True)
        
        # 生成策略报告
        policy_report = {
            "defense_strategies": []
        }
        
        total_cost = 0
        
        for i in range(self.env.num_defenses):
            if action[i] == 1:
                defense = self.env.defenses[i]
                policy_report["defense_strategies"].append({
                    "name": defense["name"],
                    "cost": defense["cost"],
                    "effectiveness": defense["effectiveness"],
                    "coverage": defense["coverage"]
                })
                total_cost += defense["cost"]
        
        policy_report["total_cost"] = total_cost
        
        # 保存策略报告
        report_path = os.path.join(self.results_dir, "security_policy.json")
        import json
        with open(report_path, 'w', encoding='utf-8') as f:
            json.dump(policy_report, f, ensure_ascii=False, indent=2)
        
        print(f"优化的安全策略已保存到: {report_path}")
        
        # 打印策略报告
        print("\n===== 优化的安全策略 =====")
        print(f"总防御成本: {total_cost}")
        print("推荐的防御策略:")
        for strategy in policy_report["defense_strategies"]:
            print(f"- {strategy['name']} (成本: {strategy['cost']})")
        
        return policy_report

# 主函数
def main():
    print("===== AI驱动的安全防护策略优化工具 =====")
    
    # 创建安全防护优化器实例
    optimizer = SecurityDefenseOptimizer()
    
    # 训练模型
    print("训练强化学习模型...")
    model = optimizer.train_model(total_timesteps=50000)
    
    # 评估模型
    print("评估模型性能...")
    results = optimizer.evaluate_model(num_episodes=20)
    
    # 生成优化的安全策略
    print("生成优化的安全策略...")
    policy = optimizer.generate_security_policy()
    
    print("\n工具使用说明：")
    print("1. 在SecurityEnv类中修改资产、攻击和防御策略的参数，以适应您的系统环境")
    print("2. 调整训练参数（如total_timesteps）以获得更好的优化效果")
    print("3. 在results目录中查看评估结果和可视化图表")
    print("4. 使用optimizer.generate_security_policy()函数生成适用于您系统的安全策略")

if __name__ == "__main__":
    main()

防护策略 → 对抗攻击 → 策略弱点发现 → 策略优化 → 再次测试 → 持续优化

数据采集 → 数据预处理 → 特征工程 → 模型训练 → 模型评估 → 漏洞检测 → 风险分析 → 防护策略生成 → 策略验证 → 报告生成

监控 → 告警 → 分析 → 优化 → 更新 → 验证 → 反馈

需求分析 → 框架设计 → 数据准备 → 模型训练 → 系统集成 → 测试验证 → 上线运行 → 持续优化

趋势1：AI技术深度融合 → 趋势2：自动化与智能化增强 → 趋势3：安全测试左移 → 趋势4：全生命周期安全保障 → 趋势5：自适应安全防护

团队组建 → 技能培养 → 流程优化 → 文化建设 → 持续改进

AI驱动安全测试价值分布: 效率提升(35%) | 准确性提高(30%) | 覆盖增强(20%) | 成本降低(15%)