模糊测试突破：AI如何优化Fuzzing效率

# 基于强化学习的智能Fuzzing工具示例
# 运行环境：Python 3.8+, pip install tensorflow numpy gym stable-baselines3

import numpy as np
import tensorflow as tf
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Dense, LSTM, Input
import gym
from gym import spaces
import time
import random

# 定义Fuzzing测试环境
class FuzzingEnv(gym.Env):
    def __init__(self, target_function, input_size=10, max_steps=100):
        super(FuzzingEnv, self).__init__()
        
        # 目标函数（被测试的函数）
        self.target_function = target_function
        # 输入大小
        self.input_size = input_size
        # 最大步数
        self.max_steps = max_steps
        # 当前步数
        self.current_step = 0
        # 当前输入
        self.current_input = None
        # 历史覆盖率
        self.history_coverage = []
        # 发现的崩溃次数
        self.crash_count = 0
        
        # 定义动作空间：修改输入的某个字节
        self.action_space = spaces.Discrete(input_size * 256)  # 每个字节可以修改为0-255
        
        # 定义观测空间：输入的字节值和覆盖率信息
        self.observation_space = spaces.Box(
            low=0,
            high=255,
            shape=(input_size + 1,),  # 输入字节 + 覆盖率
            dtype=np.uint8
        )
    
    def reset(self):
        # 重置环境状态
        self.current_step = 0
        self.current_input = np.random.randint(0, 256, size=self.input_size, dtype=np.uint8)
        self.history_coverage = []
        self.crash_count = 0
        
        # 返回初始观测值
        initial_coverage = 0.0  # 初始覆盖率为0
        obs = np.concatenate((self.current_input, [initial_coverage * 255]))  # 归一化到0-255
        return obs
    
    def step(self, action):
        # 解析动作：修改哪个字节，修改为什么值
        byte_pos = action // 256
        byte_value = action % 256
        
        # 创建新的测试输入
        new_input = self.current_input.copy()
        new_input[byte_pos] = byte_value
        
        # 执行目标函数，获取覆盖率和是否崩溃
        try:
            coverage = self.target_function(new_input)
            crashed = False
        except Exception as e:
            coverage = 0.0  # 发生崩溃时，假设覆盖率为0
            crashed = True
            self.crash_count += 1
        
        # 更新历史覆盖率
        self.history_coverage.append(coverage)
        
        # 计算奖励
        # 1. 覆盖率奖励：基于当前覆盖率和历史覆盖率的变化
        if len(self.history_coverage) > 1:
            coverage_reward = (coverage - self.history_coverage[-2]) * 10
        else:
            coverage_reward = coverage * 10
        
        # 2. 崩溃奖励：发现崩溃时给予额外奖励
        crash_reward = 100 if crashed else 0
        
        # 总奖励
        reward = coverage_reward + crash_reward
        
        # 更新当前输入
        self.current_input = new_input
        
        # 更新步数
        self.current_step += 1
        
        # 检查是否达到终止条件
        done = self.current_step >= self.max_steps
        
        # 准备观测值
        obs = np.concatenate((self.current_input, [coverage * 255]))  # 归一化到0-255
        
        # 返回观测值、奖励、是否终止、额外信息
        info = {
            "coverage": coverage,
            "crashed": crashed,
            "input": new_input
        }
        
        return obs, reward, done, info

# 构建强化学习智能体
class FuzzingAgent:
    def __init__(self, state_size, action_size):
        self.state_size = state_size
        self.action_size = action_size
        self.memory = []
        self.gamma = 0.95  # 折扣因子
        self.epsilon = 1.0  # 探索率
        self.epsilon_min = 0.01
        self.epsilon_decay = 0.995
        self.learning_rate = 0.001
        self.model = self._build_model()
    
    def _build_model(self):
        # 构建Q网络模型
        model = Sequential([
            Input(shape=(self.state_size,)),
            Dense(24, activation='relu'),
            Dense(24, activation='relu'),
            Dense(self.action_size, activation='linear')
        ])
        
        model.compile(loss='mse', optimizer=tf.keras.optimizers.Adam(learning_rate=self.learning_rate))
        return model
    
    def remember(self, state, action, reward, next_state, done):
        # 存储经验
        self.memory.append((state, action, reward, next_state, done))
    
    def act(self, state):
        # epsilon-贪婪策略选择动作
        if np.random.rand() <= self.epsilon:
            return random.randrange(self.action_size)
        act_values = self.model.predict(state.reshape(1, -1), verbose=0)
        return np.argmax(act_values[0])
    
    def replay(self, batch_size):
        # 经验回放学习
        minibatch = random.sample(self.memory, min(len(self.memory), batch_size))
        
        for state, action, reward, next_state, done in minibatch:
            target = reward
            if not done:
                target = reward + self.gamma * np.amax(self.model.predict(next_state.reshape(1, -1), verbose=0)[0])
            target_f = self.model.predict(state.reshape(1, -1), verbose=0)
            target_f[0][action] = target
            self.model.fit(state.reshape(1, -1), target_f, epochs=1, verbose=0)
        
        if self.epsilon > self.epsilon_min:
            self.epsilon *= self.epsilon_decay

# 模拟目标函数（被测试的函数）
def target_function(input_bytes):
    """
    模拟一个有漏洞的目标函数
    这里只是一个示例，实际应用中需要替换为真实的被测试函数
    """
    # 将字节数组转换为字符串（示例）
    input_str = ''.join([chr(b) for b in input_bytes])
    
    # 模拟覆盖率计算（实际应用中需要使用真实的代码覆盖率工具）
    coverage = 0.0
    
    # 检查是否包含特定模式，模拟代码覆盖
    if 'A' in input_str: coverage += 0.2
    if 'B' in input_str: coverage += 0.2
    if 'C' in input_str: coverage += 0.2
    if 'D' in input_str: coverage += 0.2
    if 'E' in input_str: coverage += 0.2
    
    # 模拟一个漏洞：如果输入包含特定序列，触发崩溃
    if 'ABCDE' in input_str:
        raise Exception("模拟缓冲区溢出")
    
    return coverage

# 主函数
def main():
    # 创建Fuzzing环境
    env = FuzzingEnv(target_function, input_size=10, max_steps=100)
    
    # 创建强化学习智能体
    state_size = env.observation_space.shape[0]
    action_size = env.action_space.n
    agent = FuzzingAgent(state_size, action_size)
    
    # 设置训练参数
    episodes = 50  # 训练回合数
    batch_size = 32  # 经验回放批次大小
    
    # 训练智能体
    print("开始训练智能Fuzzing智能体...")
    
    for e in range(episodes):
        state = env.reset()
        total_reward = 0
        
        for time_step in range(env.max_steps):
            # 智能体选择动作
            action = agent.act(state)
            
            # 执行动作，获取反馈
            next_state, reward, done, info = env.step(action)
            
            # 存储经验
            agent.remember(state, action, reward, next_state, done)
            
            # 更新状态和总奖励
            state = next_state
            total_reward += reward
            
            # 如果发现崩溃，输出信息
            if info["crashed"]:
                print(f"回合 {e+1}, 步骤 {time_step+1}: 发现崩溃! 输入: {info['input']}")
            
            # 如果达到最大步数，结束回合
            if done:
                print(f"回合 {e+1}/{episodes} 完成, 总奖励: {total_reward:.2f}, 覆盖率: {info['coverage']:.2f}, 崩溃次数: {env.crash_count}")
                break
        
        # 经验回放学习
        if len(agent.memory) > batch_size:
            agent.replay(batch_size)
    
    # 使用训练好的智能体进行Fuzzing测试
    print("\n开始使用训练好的智能体进行Fuzzing测试...")
    
    # 禁用探索，只使用 exploitation
    agent.epsilon = 0.0
    
    # 运行测试
    total_crashes = 0
    for test_run in range(10):
        state = env.reset()
        
        for time_step in range(env.max_steps):
            action = agent.act(state)
            next_state, reward, done, info = env.step(action)
            state = next_state
            
            if info["crashed"]:
                total_crashes += 1
                print(f"测试运行 {test_run+1}, 步骤 {time_step+1}: 发现崩溃! 输入: {info['input']}")
            
            if done:
                break
    
    print(f"\n测试完成，共发现 {total_crashes} 次崩溃")

if __name__ == "__main__":
    main()