AI助力CTF取证：数字痕迹的智能追踪与分析

# 示例代码：文件特征提取
import os
import hashlib
import numpy as np
from collections import Counter

# 计算文件哈希值
def calculate_file_hash(file_path, hash_type='md5'):
    hash_obj = hashlib.new(hash_type)
    with open(file_path, 'rb') as f:
        for chunk in iter(lambda: f.read(4096), b''):
            hash_obj.update(chunk)
    return hash_obj.hexdigest()

# 提取文件基本特征
def extract_file_basic_features(file_path):
    features = {}
    # 文件大小
    features['size'] = os.path.getsize(file_path)
    # 文件扩展名
    _, ext = os.path.splitext(file_path)
    features['extension'] = ext.lower() if ext else 'no_ext'
    # 文件哈希值
    features['md5'] = calculate_file_hash(file_path, 'md5')
    features['sha1'] = calculate_file_hash(file_path, 'sha1')
    return features

# 提取文件内容特征
def extract_file_content_features(file_path, max_bytes=1024):
    features = {}
    try:
        with open(file_path, 'rb') as f:
            content = f.read(max_bytes)
        
        # 字节频率统计
        byte_counter = Counter(content)
        features['byte_freq'] = {i: byte_counter.get(i, 0)/len(content) for i in range(256)}
        
        # 熵值计算（衡量数据随机性）
        entropy = 0
        for freq in features['byte_freq'].values():
            if freq > 0:
                entropy -= freq * np.log2(freq)
        features['entropy'] = entropy
        
        # 空字节比例
        features['null_byte_ratio'] = byte_counter.get(0, 0)/len(content)
        
        # 可打印字符比例
        printable_count = sum(1 for b in content if 32 <= b <= 126)
        features['printable_ratio'] = printable_count/len(content)
        
    except Exception as e:
        print(f"Error extracting features from {file_path}: {e}")
        # 设置默认值
        features['byte_freq'] = {i: 0 for i in range(256)}
        features['entropy'] = 0
        features['null_byte_ratio'] = 0
        features['printable_ratio'] = 0
    
    return features

# 提取综合文件特征
def extract_file_features(file_path):
    # 提取基本特征
    basic_features = extract_file_basic_features(file_path)
    # 提取内容特征
    content_features = extract_file_content_features(file_path)
    # 合并特征
    combined_features = {**basic_features, **content_features}
    return combined_features

# 运行环境：Python 3.8+, 需要安装 volatility3, tensorflow, numpy, matplotlib等库
import os
import numpy as np
import tensorflow as tf
from tensorflow.keras.models import Model
from tensorflow.keras.layers import Input, Conv2D, MaxPooling2D, Flatten, Dense, Dropout
import matplotlib.pyplot as plt
import volatility3
from volatility3.framework import contexts
from volatility3.framework.interfaces.plugins import PluginInterface
from volatility3.plugins.linux import pslist

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

# 加载内存映像
def load_memory_image(image_path):
    context = contexts.Context()
    try:
        # 使用volatility3加载内存映像
        # 这里仅作为示例，实际使用时需要根据内存映像的类型和格式进行相应的设置
        print(f"正在加载内存映像: {image_path}")
        # 实际代码中需要使用volatility3的API进行详细的内存解析
        # 由于volatility3的使用较为复杂，这里省略了具体的实现细节
        return context
    except Exception as e:
        print(f"加载内存映像失败: {e}")
        return None

# 提取内存段特征
def extract_memory_segment_features(segment_data, segment_size=4096):
    # 确保segment_data的大小为segment_size
    if len(segment_data) < segment_size:
        # 如果数据不足，用0填充
        padded_data = np.zeros(segment_size, dtype=np.uint8)
        padded_data[:len(segment_data)] = np.frombuffer(segment_data, dtype=np.uint8)
        segment_data = padded_data.tobytes()
    else:
        # 如果数据过多，截取前segment_size个字节
        segment_data = segment_data[:segment_size]
    
    # 将字节数据转换为2D数组（64x64）
    data_array = np.frombuffer(segment_data, dtype=np.uint8)
    data_array = data_array.reshape((64, 64))
    
    # 扩展维度以匹配CNN输入要求
    data_array = np.expand_dims(data_array, axis=-1)
    data_array = np.expand_dims(data_array, axis=0)
    
    # 归一化
    data_array = data_array / 255.0
    
    return data_array

# 创建恶意代码检测模型
def create_malware_detection_model(input_shape=(64, 64, 1)):
    # 输入层
    inputs = Input(shape=input_shape)
    
    # CNN特征提取层
    x = Conv2D(32, (3, 3), activation='relu', padding='same')(inputs)
    x = MaxPooling2D((2, 2))(x)
    x = Conv2D(64, (3, 3), activation='relu', padding='same')(x)
    x = MaxPooling2D((2, 2))(x)
    x = Conv2D(128, (3, 3), activation='relu', padding='same')(x)
    x = MaxPooling2D((2, 2))(x)
    
    # 全连接层
    x = Flatten()(x)
    x = Dense(128, activation='relu')(x)
    x = Dropout(0.5)(x)
    outputs = Dense(1, activation='sigmoid')(x)
    
    # 构建模型
    model = Model(inputs=inputs, outputs=outputs)
    
    # 编译模型
    model.compile(optimizer='adam',
                  loss='binary_crossentropy',
                  metrics=['accuracy'])
    
    return model

# 检测内存段中的恶意代码
def detect_malware_in_segment(segment_data, model):
    # 提取特征
    features = extract_memory_segment_features(segment_data)
    
    # 模型预测
    prediction = model.predict(features)
    
    # 解析预测结果
    is_malicious = prediction[0][0] > 0.5
    confidence = prediction[0][0]
    
    return is_malicious, confidence

# 可视化检测结果
def visualize_detection_results(processes, results, output_path='detection_results.png'):
    # 准备数据
    process_names = [p['name'] for p in processes]
    malicious_scores = [r['confidence'] for r in results]
    
    # 创建图表
    plt.figure(figsize=(12, 8))
    bars = plt.bar(process_names, malicious_scores, color=['red' if s > 0.5 else 'green' for s in malicious_scores])
    plt.xlabel('进程名称')
    plt.ylabel('恶意代码置信度')
    plt.title('内存进程恶意代码检测结果')
    plt.xticks(rotation=90)
    plt.axhline(y=0.5, color='gray', linestyle='--', label='阈值')
    plt.legend()
    plt.tight_layout()
    
    # 保存图表
    plt.savefig(output_path)
    print(f"检测结果图表已保存至: {output_path}")

# 主函数
def main():
    # 内存映像文件路径
    memory_image_path = 'suspicious_memory.dmp'
    
    # 加载内存映像
    context = load_memory_image(memory_image_path)
    if context is None:
        print("无法加载内存映像，程序退出。")
        return
    
    # 创建恶意代码检测模型
    model = create_malware_detection_model()
    
    # 尝试加载预训练模型权重
    try:
        model.load_weights('malware_detection_model_weights.h5')
        print("已加载预训练模型权重。")
    except Exception as e:
        print(f"无法加载预训练模型权重: {e}")
        print("将使用随机初始化的模型进行检测。")
    
    # 这里应该使用volatility3提取进程和内存段信息
    # 由于volatility3的使用较为复杂，这里使用模拟数据进行演示
    print("正在分析内存映像中的进程...")
    
    # 模拟进程数据
    simulated_processes = [
        {'name': 'explorer.exe', 'pid': 1234, 'memory_segment': os.urandom(4096)},
        {'name': 'svchost.exe', 'pid': 5678, 'memory_segment': os.urandom(4096)},
        {'name': 'suspicious.exe', 'pid': 9012, 'memory_segment': os.urandom(4096)},
        {'name': 'chrome.exe', 'pid': 3456, 'memory_segment': os.urandom(4096)},
        {'name': 'malware.exe', 'pid': 7890, 'memory_segment': os.urandom(4096)}
    ]
    
    # 检测每个进程的内存段
    detection_results = []
    for proc in simulated_processes:
        print(f"正在检测进程: {proc['name']} (PID: {proc['pid']})")
        is_malicious, confidence = detect_malware_in_segment(proc['memory_segment'], model)
        
        result = {
            'process_name': proc['name'],
            'pid': proc['pid'],
            'is_malicious': is_malicious,
            'confidence': confidence
        }
        
        detection_results.append(result)
        
        # 输出检测结果
        status = "恶意" if is_malicious else "正常"
        print(f"  检测结果: {status} (置信度: {confidence:.4f})")
    
    # 可视化检测结果
    visualize_detection_results(simulated_processes, detection_results)
    
    # 输出总结报告
    malicious_count = sum(1 for r in detection_results if r['is_malicious'])
    total_count = len(detection_results)
    print(f"\n检测总结:")
    print(f"共检测 {total_count} 个进程")
    print(f"发现 {malicious_count} 个可疑进程")
    
    if malicious_count > 0:
        print("可疑进程列表:")
        for r in detection_results:
            if r['is_malicious']:
                print(f"  - {r['process_name']} (PID: {r['pid']}, 置信度: {r['confidence']:.4f})")

if __name__ == "__main__":
    main()