AI破解CTF：漏洞挖掘与利用的智能革命

# 基于深度学习的自动化漏洞挖掘系统示例
# 运行环境：Python 3.8+, pip install tensorflow scikit-learn pandas numpy matplotlib lief

import os
import re
import lief
import numpy as np
import pandas as pd
import tensorflow as tf
from tensorflow.keras.models import Model, load_model
from tensorflow.keras.layers import Input, Dense, Conv1D, MaxPooling1D, Dropout, LSTM, Bidirectional, Attention, Concatenate
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler, LabelEncoder
from sklearn.metrics import classification_report, confusion_matrix, accuracy_score
import matplotlib.pyplot as plt
import seaborn as sns
from concurrent.futures import ThreadPoolExecutor
import time
import tempfile
import subprocess
from typing import Dict, List, Tuple, Any, Optional

# 设置中文显示
plt.rcParams["font.family"] = ["SimHei", "WenQuanYi Micro Hei", "Heiti TC"]
plt.rcParams["axes.unicode_minus"] = False  # 解决负号显示问题

# 设置随机种子，确保结果可复现
def set_seed(seed=42):
    np.random.seed(seed)
    tf.random.set_seed(seed)
    os.environ["PYTHONHASHSEED"] = str(seed)

def get_device():
    """获取可用的计算设备"""
    if tf.config.list_physical_devices('GPU'):
        return "GPU"
    else:
        return "CPU"

def disassemble_binary(file_path: str) -> List[Dict]:
    """反汇编二进制文件，提取指令序列"""
    try:
        binary = lief.parse(file_path)
        instructions = []
        
        # 遍历所有函数
        for func in binary.functions:
            func_name = func.name if func.name else f"func_{hex(func.address)}"
            func_instructions = []
            
            # 提取函数的指令
            for instr in func.instructions:
                instr_info = {
                    "address": hex(instr.address),
                    "mnemonic": instr.mnemonic,
                    "operands": str(instr.operands),
                    "bytes": instr.bytes
                }
                func_instructions.append(instr_info)
            
            instructions.append({
                "name": func_name,
                "address": hex(func.address),
                "instructions": func_instructions
            })
        
        return instructions
    except Exception as e:
        print(f"反汇编二进制文件失败: {e}")
        return []

def extract_features(instructions: List[Dict]) -> np.ndarray:
    """从指令序列中提取特征"""
    # 简单的特征提取示例，实际应用中可能需要更复杂的特征工程
    features = []
    
    # 定义感兴趣的指令类型
    memory_instructions = set(["mov", "push", "pop", "lea", "cmp", "test", "add", "sub", "inc", "dec"])
    arithmetic_instructions = set(["add", "sub", "mul", "div", "inc", "dec", "neg", "xor", "and", "or"])
    control_flow_instructions = set(["jmp", "jz", "jnz", "je", "jne", "jb", "jnb", "ja", "jna", "call", "ret"])
    
    for func in instructions:
        func_features = []
        func_instrs = func["instructions"]
        
        # 基本统计特征
        total_instructions = len(func_instrs)
        memory_instr_count = sum(1 for instr in func_instrs if instr["mnemonic"].lower() in memory_instructions)
        arithmetic_instr_count = sum(1 for instr in func_instrs if instr["mnemonic"].lower() in arithmetic_instructions)
        control_flow_instr_count = sum(1 for instr in func_instrs if instr["mnemonic"].lower() in control_flow_instructions)
        
        # 计算比例特征
        memory_ratio = memory_instr_count / total_instructions if total_instructions > 0 else 0
        arithmetic_ratio = arithmetic_instr_count / total_instructions if total_instructions > 0 else 0
        control_flow_ratio = control_flow_instr_count / total_instructions if total_instructions > 0 else 0
        
        # 检查是否包含可能存在漏洞的指令序列
        has_strcpy = any("strcpy" in instr["operands"].lower() for instr in func_instrs)
        has_memcpy = any("memcpy" in instr["operands"].lower() for instr in func_instrs)
        has_printf = any("printf" in instr["operands"].lower() for instr in func_instrs) and any("%" in instr["operands"] for instr in func_instrs)
        has_gets = any("gets" in instr["operands"].lower() for instr in func_instrs)
        
        # 提取指令的操作数特征
        immediate_operands = sum(1 for instr in func_instrs if any(c.isdigit() for c in instr["operands"]))
        register_operands = sum(1 for instr in func_instrs if any(reg in instr["operands"].lower() for reg in ["eax", "ebx", "ecx", "edx", "esp", "ebp"]))
        memory_operands = sum(1 for instr in func_instrs if any(mem in instr["operands"].lower() for mem in ["[eax]", "[ebx]", "[ecx]", "[edx]", "[esp]", "[ebp]"]))
        
        # 构建特征向量
        func_features.extend([
            total_instructions,
            memory_ratio,
            arithmetic_ratio,
            control_flow_ratio,
            immediate_operands / total_instructions if total_instructions > 0 else 0,
            register_operands / total_instructions if total_instructions > 0 else 0,
            memory_operands / total_instructions if total_instructions > 0 else 0,
            1 if has_strcpy else 0,
            1 if has_memcpy else 0,
            1 if has_printf else 0,
            1 if has_gets else 0
        ])
        
        features.append(func_features)
    
    return np.array(features)

class VulnerabilityDetector:
    """基于深度学习的漏洞检测器"""
    def __init__(self):
        self.model = None
        self.scaler = StandardScaler()
        self.label_encoder = LabelEncoder()
        self.input_shape = None
        
    def build_model(self, input_shape: Tuple, num_classes: int = 2):
        """构建深度学习模型"""
        inputs = Input(shape=input_shape)
        
        # CNN部分用于提取局部特征
        cnn_output = Conv1D(64, kernel_size=3, activation='relu')(inputs)
        cnn_output = MaxPooling1D(pool_size=2)(cnn_output)
        cnn_output = Dropout(0.3)(cnn_output)
        cnn_output = Conv1D(128, kernel_size=3, activation='relu')(cnn_output)
        cnn_output = MaxPooling1D(pool_size=2)(cnn_output)
        cnn_output = Dropout(0.3)(cnn_output)
        
        # LSTM部分用于提取序列特征
        lstm_output = Bidirectional(LSTM(64, return_sequences=True))(inputs)
        lstm_output = Bidirectional(LSTM(128))(lstm_output)
        
        # 融合CNN和LSTM的输出
        merged = Concatenate()([cnn_output, lstm_output])
        
        # 全连接层进行分类
        x = Dense(128, activation='relu')(merged)
        x = Dropout(0.5)(x)
        x = Dense(64, activation='relu')(x)
        
        # 输出层
        if num_classes == 2:
            outputs = Dense(1, activation='sigmoid')(x)
            self.model = Model(inputs=inputs, outputs=outputs)
            self.model.compile(optimizer='adam', loss='binary_crossentropy', metrics=['accuracy'])
        else:
            outputs = Dense(num_classes, activation='softmax')(x)
            self.model = Model(inputs=inputs, outputs=outputs)
            self.model.compile(optimizer='adam', loss='sparse_categorical_crossentropy', metrics=['accuracy'])
        
        self.input_shape = input_shape
        
        # 打印模型结构
        self.model.summary()
    
    def train(self, X_train: np.ndarray, y_train: np.ndarray, X_val: np.ndarray = None, y_val: np.ndarray = None, 
              epochs: int = 50, batch_size: int = 32):
        """训练模型"""
        if self.model is None:
            raise ValueError("模型尚未构建，请先调用build_model方法")
        
        # 数据预处理
        X_train_scaled = self.scaler.fit_transform(X_train)
        
        # 确保输入形状正确
        if len(X_train_scaled.shape) == 2:
            X_train_scaled = X_train_scaled.reshape(X_train_scaled.shape[0], X_train_scaled.shape[1], 1)
        
        # 准备验证集
        validation_data = None
        if X_val is not None and y_val is not None:
            X_val_scaled = self.scaler.transform(X_val)
            if len(X_val_scaled.shape) == 2:
                X_val_scaled = X_val_scaled.reshape(X_val_scaled.shape[0], X_val_scaled.shape[1], 1)
            validation_data = (X_val_scaled, y_val)
        
        # 训练模型
        history = self.model.fit(
            X_train_scaled, y_train,
            validation_data=validation_data,
            epochs=epochs,
            batch_size=batch_size,
            callbacks=[
                tf.keras.callbacks.EarlyStopping(monitor='val_loss', patience=10, restore_best_weights=True),
                tf.keras.callbacks.ReduceLROnPlateau(monitor='val_loss', factor=0.5, patience=5)
            ]
        )
        
        return history
    
    def evaluate(self, X_test: np.ndarray, y_test: np.ndarray):
        """评估模型性能"""
        if self.model is None:
            raise ValueError("模型尚未构建，请先调用build_model方法")
        
        # 数据预处理
        X_test_scaled = self.scaler.transform(X_test)
        if len(X_test_scaled.shape) == 2:
            X_test_scaled = X_test_scaled.reshape(X_test_scaled.shape[0], X_test_scaled.shape[1], 1)
        
        # 评估模型
        loss, accuracy = self.model.evaluate(X_test_scaled, y_test, verbose=1)
        
        # 进行预测
        if self.model.output_shape[-1] == 1:
            y_pred_prob = self.model.predict(X_test_scaled)
            y_pred = (y_pred_prob > 0.5).astype(int)
        else:
            y_pred = np.argmax(self.model.predict(X_test_scaled), axis=1)
        
        # 生成分类报告
        report = classification_report(y_test, y_pred)
        cm = confusion_matrix(y_test, y_pred)
        
        return {
            "loss": loss,
            "accuracy": accuracy,
            "classification_report": report,
            "confusion_matrix": cm
        }
    
    def predict(self, X: np.ndarray):
        """使用模型进行预测"""
        if self.model is None:
            raise ValueError("模型尚未构建，请先调用build_model方法")
        
        # 数据预处理
        X_scaled = self.scaler.transform(X)
        if len(X_scaled.shape) == 2:
            X_scaled = X_scaled.reshape(X_scaled.shape[0], X_scaled.shape[1], 1)
        
        # 进行预测
        if self.model.output_shape[-1] == 1:
            y_pred_prob = self.model.predict(X_scaled)
            y_pred = (y_pred_prob > 0.5).astype(int)
            return y_pred, y_pred_prob
        else:
            y_pred_prob = self.model.predict(X_scaled)
            y_pred = np.argmax(y_pred_prob, axis=1)
            return y_pred, y_pred_prob
    
    def save_model(self, model_path: str):
        """保存模型"""
        if self.model is None:
            raise ValueError("模型尚未构建，请先调用build_model方法")
        
        # 创建保存目录
        os.makedirs(os.path.dirname(model_path), exist_ok=True)
        
        # 保存模型
        self.model.save(model_path)
        
        # 保存scaler
        np.save(f"{model_path}_scaler_mean.npy", self.scaler.mean_)
        np.save(f"{model_path}_scaler_var.npy", self.scaler.var_)
    
    def load_model(self, model_path: str):
        """加载模型"""
        # 加载模型
        self.model = load_model(model_path)
        
        # 加载scaler
        mean = np.load(f"{model_path}_scaler_mean.npy")
        var = np.load(f"{model_path}_scaler_var.npy")
        self.scaler.mean_ = mean
        self.scaler.var_ = var
        self.scaler.scale_ = np.sqrt(var)
        
        # 获取输入形状
        self.input_shape = self.model.input_shape[1:]

class AutomatedVulnerabilityScanner:
    """自动化漏洞扫描器"""
    def __init__(self):
        self.detector = VulnerabilityDetector()
        self.binary_path = None
        self.functions = []
        self.features = None
        self.results = {}
        
    def load_binary(self, file_path: str) -> bool:
        """加载二进制文件"""
        try:
            self.binary_path = file_path
            self.functions = disassemble_binary(file_path)
            return True
        except Exception as e:
            print(f"加载二进制文件失败: {e}")
            return False
    
    def extract_features(self):
        """提取特征"""
        if not self.functions:
            print("没有可提取特征的函数，请先加载二进制文件")
            return False
        
        try:
            self.features = extract_features(self.functions)
            return True
        except Exception as e:
            print(f"提取特征失败: {e}")
            return False
    
    def load_pretrained_model(self, model_path: str) -> bool:
        """加载预训练模型"""
        try:
            self.detector.load_model(model_path)
            return True
        except Exception as e:
            print(f"加载预训练模型失败: {e}")
            return False
    
    def scan(self) -> Dict:
        """扫描漏洞"""
        if self.features is None:
            print("没有可扫描的特征，请先提取特征")
            return {}
        
        try:
            # 进行预测
            y_pred, y_pred_prob = self.detector.predict(self.features)
            
            # 存储结果
            vulnerability_results = []
            for i, func in enumerate(self.functions):
                func_name = func["name"]
                func_address = func["address"]
                is_vulnerable = bool(y_pred[i])
                confidence = float(y_pred_prob[i]) if len(y_pred_prob.shape) == 1 else float(y_pred_prob[i][y_pred[i]])
                
                vulnerability_results.append({
                    "name": func_name,
                    "address": func_address,
                    "is_vulnerable": is_vulnerable,
                    "confidence": confidence
                })
            
            self.results = {
                "total_functions": len(self.functions),
                "vulnerable_functions": sum(y_pred),
                "results": vulnerability_results
            }
            
            return self.results
        except Exception as e:
            print(f"扫描漏洞失败: {e}")
            return {}
    
    def generate_report(self, report_path: str = None):
        """生成漏洞扫描报告"""
        if not self.results:
            print("没有可生成报告的结果，请先进行扫描")
            return False
        
        try:
            # 创建报告内容
            report = "# 漏洞扫描报告\n\n"
            report += f"## 扫描摘要\n"
            report += f"- 扫描文件: {self.binary_path}\n"
            report += f"- 扫描时间: {time.strftime('%Y-%m-%d %H:%M:%S')}\n"
            report += f"- 总函数数: {self.results['total_functions']}\n"
            report += f"- 潜在漏洞函数数: {self.results['vulnerable_functions']}\n\n"
            
            if self.results['vulnerable_functions'] > 0:
                report += "## 潜在漏洞函数列表\n"
                report += "| 函数名称 | 函数地址 | 置信度 |\n"
                report += "|---------|---------|-------|\n"
                
                # 按置信度排序
                vulnerable_funcs = sorted(
                    [f for f in self.results['results'] if f['is_vulnerable']],
                    key=lambda x: x['confidence'],
                    reverse=True
                )
                
                for func in vulnerable_funcs:
                    report += f"| {func['name']} | {func['address']} | {func['confidence']:.4f} |\n"
            else:
                report += "## 潜在漏洞函数列表\n"
                report += "未发现潜在漏洞函数。\n"
            
            # 保存报告
            if report_path:
                os.makedirs(os.path.dirname(report_path), exist_ok=True)
                with open(report_path, 'w', encoding='utf-8') as f:
                    f.write(report)
                print(f"报告已保存至: {report_path}")
            
            return report
        except Exception as e:
            print(f"生成报告失败: {e}")
            return False

# 主函数
def main():
    print("开始基于深度学习的自动化漏洞挖掘系统演示...")
    print(f"使用设备: {get_device()}")
    
    # 创建自动化漏洞扫描器
    scanner = AutomatedVulnerabilityScanner()
    
    # 在实际环境中，你可能需要提供二进制文件的路径和预训练模型的路径
    binary_path = input("请输入要扫描的二进制文件路径: ")
    
    # 注意：在实际应用中，你需要有一个预训练的模型
    # 这里为了演示，我们假设模型已经训练好并保存在指定路径
    # model_path = "models/vulnerability_detector.h5"
    
    # 加载二进制文件
    print(f"\n1. 加载二进制文件: {binary_path}")
    if not scanner.load_binary(binary_path):
        print("加载二进制文件失败，退出程序。")
        return
    
    # 提取特征
    print("\n2. 提取函数特征...")
    if not scanner.extract_features():
        print("提取特征失败，退出程序。")
        return
    
    print(f"成功提取 {scanner.features.shape[0]} 个函数的特征。")
    
    # 在实际应用中，你需要加载预训练模型
    # 由于是演示，这里跳过了加载模型和扫描的步骤
    # 下面是实际应用中的代码示例
    
    # print("\n3. 加载预训练模型...")
    # if not scanner.load_pretrained_model(model_path):
    #     print("加载预训练模型失败，退出程序。")
    #     return
    
    # print("\n4. 扫描漏洞...")
    # results = scanner.scan()
    # print(f"扫描完成，发现 {results['vulnerable_functions']} 个潜在漏洞函数。")
    
    # print("\n5. 生成漏洞扫描报告...")
    # report = scanner.generate_report("reports/vulnerability_report.md")
    # print("漏洞扫描报告预览:")
    # print(report)
    
    print("\n注意：由于这是一个演示程序，我们没有实际的预训练模型用于漏洞检测。")
    print("在实际应用中，你需要先收集漏洞样本数据，训练模型，然后使用训练好的模型进行漏洞扫描。")
    print("\n基于深度学习的自动化漏洞挖掘系统演示完成！")

if __name__ == "__main__":
    set_seed()
    main()