民间机器学习反病毒技术发展回顾
--一种基于DRL(Deep Reinforcement Learning)的恶意程序检测设想
# 引言
随着计算机技术与机器学习技术的飞速发展,机器学习开始在反病毒领域崭露头角,同时个人开发者训练反病毒模型成为可能。
机器学习通过分析特征,实现了从特征库至泛化的飞跃。本文将分析目前社区内常见的特征工程,并提出一种基于深度强化学习(DRL)的恶意样本检测设想。
# 传统特征工程
目前社区内的特征工程主要分为两大类:
1.端-人-端式:
最传统的特征工程。通过人为划分一定段并计算数学统计量,分析PE结构等,让模型拟合人类认为相对具有区分力的特征。
优点:极高的数据利用率、相对透明可解释。
缺点:鲁棒性相对较差,需要大量样本才能得到相对可观的泛化与误报控制。
常见的手法有:IAT、可打印字符、熵、字节图等
2.端-端式:
我个人较为看好的一类特征工程。人类开发者通过适当的方法,将原始字节流直接作为特征训练。
这种方法具有较高的鲁棒性,同时维度较高,因此需要下采样才能将训练&&推理时间控制在合理范围内。
目前一个毫秒级的模型(若为神经网络)最大规模在6000维以下的输入、不超过10层的Hidden Layers。
常见的手法及个人看法:
-灰度图类:将文件全部或重要部分转换为灰度图像形式(一般不太需要下采样)进行训练。这种输入基本锁死conv2d类神经网络,如CNN。本意是在转化为纹理后找到通用的家族特征/模式。
这一手法至今仍然是广谱家族检测最有效的手段。缺点是模型推理负载一般较大,较难观察到远联系。
-N-Grams类:将文件字节流转换为机器学习可分析的形式(如uint16数字)。几乎适合所有算法。
由于信息密度极低,为了防止维度风暴,需要通过映射关系转化为N-Grams形式。如:
00 00 -> 0
00 01 -> 1
FF FF -> 65535
这样可以在不损失原始信息的情况下降低维度。几乎完全不占用任何性能。
但尤其需要注意的是,随着N的增长,单个特征的表达范围会指数级增长(16**2n),导致较近的原始信息在合并后差异较大。
如F000AF与AA00AF的压缩后文本有极大差异。这种差异过大时可能会导致算法学习到错误的模式。一般最多使用3-Grams就可以了。
-统计数据下采样:这种下采样常见于Torch类训练中。一般是取std,min,max,avg或更多的统计值。
这种下采样具有极大的信息损失,失去了语义信息,极其不推荐。更好的做法是PBC,或符号回归。
这种下采样**极易导致TabNet算法跑崩**
目前这两种类型的特征工程,要么鲁棒性差,要么信息利用率低。在规模相对小的训练集很难跑出优秀的效果。
为了解决这个问题,笔者提出了一种基于DRL的新型"特征工程"设想。
# Ti Security Solar
想要达到最好的检测效果,就应该将特征选取权下放至AI。
下面是Ti Security Solar引擎的基本架构。
-选择一个适合混合任务、长上下文的算法。这里首推PPO(Proximal Policy Optimization)。
-提取基本信息。如:文件大小、入口点偏移、文件类型(可执行文件/动态链接库/驱动程序)、导入DLL、各个节的信息(名称、大小、偏移量)
-构建特殊调用。可以当成LLM MCP工具去写。下面是Solar的一些内置工具。
0.报告此文件为恶意(True)/良性(False)。
函数原型:report(bool)
消耗积分:0
该工具仅能调用一次。调用后立刻结束此文件分析。
1.直接读取字节流。
函数原型:read_bytes(offset,size)
从offset开始,读取size Bytes的原始字节数据,并返回HEX形式。
消耗积分:size/100一次。
2.生成灰度图。
函数原型:greyscale(offset,size)
从offset开始,读取size Bytes的原始字节数据,并转换为灰度图形式返回。
消耗积分:(size/file_size)/30一次。
3.熵计算。
函数原型:low_entropy(offset,size)
从offset开始,读取size Bytes的原始字节数据,计算熵并返回。
消耗积分:[5+1.5**(n//100)]一次
4.统计数据。
函数原型:analyze_bytes(offset,size)
从offset开始,读取size Bytes的原始数据,按照出现频率百分比构建一个Dictionary。
如:["00":63.7,"AA":3.27,"9F":2.99......]
消耗积分:size/450一次
5.分层熵。
函数原型:high_entropy(offset,size,n)
从offset开始,读取size Bytes的原始字节数据,分为n层细度计算熵并返回。
如size=100,n=3,则得到:
熵[0~100],[0~50],[51~100],[0~33],[34~66],[67~100]
共1+2+3+...+n层数据
消耗积分:(5+2**n)积分/次
6.跳转查询。
函数原型:analyze_jmp(offset)
查询一条jmp指令跳转到哪里,返回相对与绝对偏移。
消耗积分:
第一次调用20积分(因为需要调用Capstone反编译。若上下文中没有调用过反编译。)
之后调用2积分/次
7.反编译结果查询。
函数原型:show_dec(offset,num)
从offset开始,返回num条指令的反编译形式(保留操作数)
第一次调用20积分(因为需要调用Capstone反编译。若上下文中没有调用过反编译。)
之后调用2积分/次
-并辅以一些DRL经典技巧:
1.初始为模型分配100点积分。积分归零未找到正确答案则积分归零,给予1单位的惩罚。
2.模型最少需要10 Steps,最多50 Steps。若n<10则视作失败,给予0.8单位惩罚。若50<n,则每多1 Step,给予模型0.02 * n的惩罚(判断错误或正确都是在原惩罚/奖励上去扣除)
3.上下文中出现完全相同的命令时,积分归零。
4.积分归零则保留必要状态重新训练该样本。
5.若判断正确,则给予0.3+剩余积分/300的奖励
6.训练时随机从模型中抽取文件。在训练1000个文件后,每训练200个文件,则在整个训练集随机挑出30个文件作测试评估。
7.开启0.2的神经元无效。
8.在训练完毕后,则保存为onnx。
9.需要强制规定模型的输出格式。
整个训练流程为:
1.脚本统计文件,safe文件夹为良性文件,virus文件夹为恶意文件。遍历。
2.开始训练。
3.训练一个文件时先mmap文件,之后将基本信息作为初始输入交给模型。之后根据模型的工具调用来返回数据。
4.一个文件必须判断正确才能训练下一个。理论上的训练速度初始很低,训练一部分学到基本模式后会开始极速加快。
# 结语
这一设想意在在最小的训练集得到最大的泛化能力。也许这才是真正的NGAV(Next-Generation Antivirus)
传统机器学习、深度学习已经基本到达能力上限,我们应何去何从?
(后补加:)
总结一下Solar的实际体验,我总结出几点
·训练只吃CPU而非GPU,这意味着训练偏慢,但金钱成本很低。
(实测40HiddenSize、10Steps,在intel Xeon Plantinum 8369B 4c8t能达到平均3文件/s)
·可以动态调节扫描深度
·推理扫描真的很快很快,瓶颈都快到IO这边了
|