AI 模糊测试：用生成模型驱动更智能的漏洞挖掘

bash pip install torch numpy matplotlib lief

我们使用 PyTorch 构建模型，LIEF 用于选择性验证 ELF 头部的合法性（可省略用 mock 判断）。

### 4.2 数据预处理与数据集
假定从系统收集一批合法 ELF 文件，提取文件头前 64 字节作为训练样本。
```python
import torch
from torch.utils.data import Dataset, DataLoader

class HeaderDataset(Dataset):
    def __init__(self, file_list, length=64):
        self.data = []
        for f in file_list:
            with open(f, 'rb') as fp:
                raw = fp.read(length)
                # 填充或截断至固定长度
                if len(raw) < length:
                    raw += b'\x00' * (length - len(raw))
                self.data.append(torch.tensor(list(raw), dtype=torch.float32))
    def __len__(self):
        return len(self.data)
    def __getitem__(self, idx):
        return self.data[idx]

真实场景需要归一化到 [0,1] 或使用 one-hot 编码。

4.3 定义 VAE 模型

import torch.nn as nn

class VAE(nn.Module):
    def __init__(self, input_dim=64, latent_dim=16):
        super().__init__()
        self.enc_fc1 = nn.Linear(input_dim, 32)
        self.enc_mu = nn.Linear(32, latent_dim)
        self.enc_logvar = nn.Linear(32, latent_dim)
        
        self.dec_fc1 = nn.Linear(latent_dim, 32)
        self.dec_out = nn.Linear(32, input_dim)
        
    def encode(self, x):
        h = torch.relu(self.enc_fc1(x))
        return self.enc_mu(h), self.enc_logvar(h)
    
    def reparameterize(self, mu, logvar):
        std = torch.exp(0.5 * logvar)
        eps = torch.randn_like(std)
        return mu + eps * std
    
    def decode(self, z):
        h = torch.relu(self.dec_fc1(z))
        return torch.sigmoid(self.dec_out(h))  # 输出归一化到 [0,1]
    
    def forward(self, x):
        mu, logvar = self.encode(x)
        z = self.reparameterize(mu, logvar)
        recon = self.decode(z)
        return recon, mu, logvar

损失函数为重建损失（MSE 或二元交叉熵）加上 KL 散度：

def loss_function(recon_x, x, mu, logvar):
    BCE = nn.functional.binary_cross_entropy(recon_x, x, reduction='sum')
    KLD = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp())
    return BCE + KLD

4.4 训练与潜在空间变异

训练完成后，我们就有了一个输入 → 潜在向量 → 重采样的映射。变异发生在潜在空间：

def mutate_from_seed(seed_bytes, model, scale=0.5):
    model.eval()
    with torch.no_grad():
        seed_tensor = torch.tensor(list(seed_bytes), dtype=torch.float32)
        mu, _ = model.encode(seed_tensor.unsqueeze(0))
        z = mu.squeeze(0)  # 使用均值作为初始潜在向量
        # 添加高斯噪声进行变异
        z_mutated = z + scale * torch.randn_like(z)
        mutated_tensor = model.decode(z_mutated.unsqueeze(0)).squeeze(0)
    # 将浮点值映射回字节
    mutated_bytes = bytes(mutated_tensor.round().clamp(0,255).int().tolist())
    return mutated_bytes