文章总结: 本文介绍基于GPT架构的webshell识别模型,参数量约10M。预训练效果不佳,SFT后出现误报,通过注入增强改进。扩大预训练规模反而增加误报,最终结论是数据量不足,模型能力未达上限。建议收集更多高质量黑样本。 综合评分: 84 文章分类: AI安全,安全工具,WEB安全,恶意软件,漏洞分析
基于GPT架构的webshell识别模型
原创
k k
漏洞推送
2026年7月11日 22:57 浙江
在小说阅读器读本章
去阅读
模型架构
基于gpt3 only decode架构
完成从头开始训练,模型总参数量: 9.98 百万,在mac m1上使用mps进行训练
模型代码
import torchimport torch.nn as nnfrom tokenizers import Tokenizerimport configimport os
class GPTBlock(nn.Module): def __init__(self, embed_size, num_heads): super().__init__() self.ln1 = nn.LayerNorm(embed_size) self.attention = nn.MultiheadAttention( embed_dim=embed_size, num_heads=num_heads, batch_first=True ) self.ln2 = nn.LayerNorm(embed_size) self.ffn = nn.Sequential( nn.Linear(embed_size, embed_size * config.FFN_SIZE), # 升维 nn.GELU(), # 非线性激活 (比 ReLU 更平滑) nn.Linear(embed_size * config.FFN_SIZE, embed_size), # 降维 )
def forward(self, x, mask=None): # 1. 注意力机制与残差连接 # PyTorch 的 MultiheadAttention 会返回 (输出, 注意力权重),我们目前只需要输出 # is_causal=True 搭配 mask 会自动处理防止“看未来”的逻辑 attn_output, _ = self.attention( x, x, x, attn_mask=mask, is_causal=True, need_weights=False, ) x = self.ln1(x + attn_output)
# 2. FFN 与残差连接 ffn_output = self.ffn(x) x = self.ln2(x + ffn_output)
return x
class SimpleGPT(nn.Module): def __init__(self, vocab_size, embed_size, max_seq_length, num_heads, num_layers): super().__init__() # 1. 词嵌入与位置嵌入 (PyTorch 的 nn.Embedding 自带可学习权重) self.token_embedding = nn.Embedding(vocab_size, embed_size) self.position_embedding = nn.Embedding(max_seq_length, embed_size)
# 2. 堆叠多个 Transformer Block (盖楼) # 用 nn.ModuleList 包装起来,确保 PyTorch 能够追踪里面的参数 self.blocks = nn.ModuleList( [GPTBlock(embed_size, num_heads) for _ in range(num_layers)] )
# 3. 语言模型头 (LM Head) self.ln_final = nn.LayerNorm(embed_size) # 投射回词表大小,bias=False 是行业惯例,节约参数 self.lm_head = nn.Linear(embed_size, vocab_size, bias=False)
def forward(self, input_ids): batch_size, seq_length = input_ids.size() device = input_ids.device
# 生成位置索引 [0, 1, ..., seq_length - 1] positions = torch.arange( 0, seq_length, dtype=torch.long, device=device ).unsqueeze(0)
# 特征融合:词意 + 位置 x = self.token_embedding(input_ids) + self.position_embedding(positions)
# 生成因果掩码 (Causal Mask) mask = nn.Transformer.generate_square_subsequent_mask(seq_length, device=device)
# 让数据一层一层穿过所有的 Transformer Block for block in self.blocks: x = block(x, mask=mask)
# 最后一层归一化 x = self.ln_final(x)
# 通过 LM Head 转化为词表打分 (Logits) # 形状变为: (Batch_Size, Seq_Len, Vocab_Size) logits = self.lm_head(x)
return logits
关键参数如下:
EMBED_SIZE = 256 # 256维特征MAX_LEN = 1024 * 4 # 4k上下文VOCAB_SIZE = 1024 * 8NUM_HEADS = 8 # 8个注意力头 (256/8 = 每个头32维)NUM_LAYERS = 6 # 堆叠6层 Transformer BlockFFN_SIZE = 4
生成词表
trainer = BpeTrainer( vocab_size=config.VOCAB_SIZE, special_tokens=["[UNK]", "[PAD]", "[BOS]", "[EOS]", "[black]", "[white]"],)
预训练
使用数据 WordPress + dede + webshell
PHP 文件数量: 2445原始 token 总量: 16,139,532当前训练实际 token 总量: 4,480,725被 MAX_LEN 截断的文件数量: 593平均原始 token/文件: 6,601.04平均训练 token/文件: 1,832.61
预训练能力测试
prompt = """<?php eval($_POST['"""
学到了一堆没屌用的
SFT
[bos_id] + code_ids + [eos_id]
在[eos_id]后增加[black]或者[white] 特殊token来标识样本
sft白数据:
sft黑数据:
•https://github.com/tennc/webshell•https://github.com/xl7dev/WebShell•https://github.com/JohnTroony/php-webshells
第一轮训练完以后出现了很严重的误报问题,会把<?php echo $admin; ?>也判断为webshell
大概原因是因为在sft数据中,黑样本的数据长度总是比较短,模型学歪了把长度作为特征了。
后续做了注入增强,1.随机抽取白样本中的片段,作为短白样本 2. 将黑样本去掉<?php ?>标签复制到白样本的;后做黑样本,并且将比例对齐
🔍 SFT 数据集加载完毕:原始 Webshell 780 个,注入增强 Webshell 780 个,正常样本 780 个,短正常样本 780 个 (原始正常 2098,原始 Webshell 780)
Eval
扩大预训练规模
接下来我想探索一个我一直很好奇的问题,在sft完全不变的情况下,增加预训练的规模对模型的最终效果是否有改善。
根据经验,模型的最近预训练token数量是模型的参数*20
这是1千万的模型,所以最佳的训练token是 2e
我在compose上收集了大概1.8e的php token,单论实际训了的token大概是1e左右。训练2论,大概在2e左右。
首先在我的mac上训练,大概训练了5-8小时,才训了是1/20。完整训练完大概要几天到一周。
不得已,租了一个5090的服务器,大概训练2个多小时即可
预训练完第一个batch,看看效果
第二个batch的效果
有比较严重的重复的问题。
sft以后的误报率也大幅升高。
推测可能是因为我在预训练数据中放了webshell的代码,webshell代码中有大量的重复,删除掉webshell代码以后再训练
Batch1
batch2
基于batch1 sft结果
基于batch2 sft结果
如果完全不做预训练的情况下进行SFT对齐
sft batch1
sft batch2
sft batch3
根据样本表现的情况,我个人觉得在模型参数量比较小的时候,模型即使经过预训练也不能完全掌握php的语法信息,导致对后期的webshell识别帮助并不大。后续训练直接做SFT训练,不再做预训练
参数对结果的影响
通过一些工程化手段,收集了GitHub上全部的webshell仓库和扩大了白样本的数量。将误报率和漏报率降下来了。
基准 10M参数
embed_size = 256
30M参数 embed_size = 512
在扩大embed_size的情况下,参数增加了20m,但是对于结果并没有显著改变。
尝试增加 FFN_SIZE = 16
20M 参数
尝试增加 FFN_SIZE = 32
30M 参数
可以见到webshell识别能力并没有明显提高,虽然FFN的siez大了很多。
我觉得可能是因为黑样本的数量还是太少,就算把GitHub上全部的黑样本收集起来去重以后也就不到3000,而且还有小部分不是webshell的噪音混在里面。而高质量黑样本数量成本很高,像阿里的伏魔比赛大概是以500每个以上的价格来收集黑样本。
因为黑样本数据太少,我个人感觉目前的数据应该是还没有触及到模型能力的上限。在数据量较少的情况下很难做到误报率和漏报率的平衡,漏报率的降低必然会显著增加误报率。
最后我把我的模型部署到我的vps上了大家可以体验一下这个
webshellgpt: https://webshell.hx-std.com/
免责声明:
本文所载程序、技术方法仅面向合法合规的安全研究与教学场景,旨在提升网络安全防护能力,具有明确的技术研究属性。
任何单位或个人未经授权,将本文内容用于攻击、破坏等非法用途的,由此引发的全部法律责任、民事赔偿及连带责任,均由行为人独立承担,本站不承担任何连带责任。
本站内容均为技术交流与知识分享目的发布,若存在版权侵权或其他异议,请通过邮件联系处理,具体联系方式可点击页面上方的联系我。
本文转载自:漏洞推送 k k《基于GPT架构的webshell识别模型》
版权声明
本站仅做备份收录,仅供研究与教学参考之用。
读者将信息用于其他用途的,全部法律及连带责任由读者自行承担,本站不承担任何责任。










评论