Gym-Malware基于OpenAI Gym和Keras-rl开发,主要由DQNAgent,MalwareEnv,interface,MalwareManipulator和PEFeatureExtractor组成。如图10-3所示,PEFeature Extractor将PE文件转换成特征向量,interface加载已经训练好的GBDT模型用于恶意程序检测,PE文件的特征向量作为状态传递。DQNAgent基于当前状态和一定的策略,选择免杀动作。MalwareEnv根据免杀动作,通过MalwareManipulator针对PE文件执行免杀操作,然后使用PEFeatureExtractor重新计算特征,再使用interface判断,如果不是恶意程序,反馈10并结束本轮学习;如果是恶意程序,反馈0以及新状态给DQNAgent,DQNAgent继续选择下一步免杀操作,如此循环。下面我们将介绍每个组件的具体原理和实现。
图10-3 Gym-Malware架构
PEFeatureExtractor的主要功能就是把PE文件转换成特征向量,关于PE文件的特征提取方法请参考本书第8章“恶意程序检测”相关内容。
PEFeatureExtractor以类的形式存在,提供了唯一的对外接口,用于把保存了PE文件的字节数组转换成特征:
PEFeatureExtractor.extract(self.bytez)
PE文件的特征分为两类:一类是通过PE文件可以直接获取到的,比如字节直方图、字节熵直方图和字符串特征;另外一类特征是需要解析PE文件结构,从各个节分析出的特征,比如节头特征、导入和导出表特征、文件头特征等。
我们定义变量raw_features,保存获取字节直方图、字节熵直方图和字符串特征的函数,代码如下:
self.raw_features = [ ByteHistogram(), ByteEntropyHistogram(), StringExtractor() ]
我们定义变量parsed_features,保存获取节头特征、导入和导出表特征、文件头特征等的函数,代码如下:
self.parsed_features = [ GeneralFileInfo(), HeaderFileInfo(), SectionInfo(), ImportsInfo(), ExportsInfo() ]
PEFeatureExtractor在解析的过程中,会调用内部函数提取特征,首先会调用raw_features中的函数提取字节直方图、字节熵直方图和字符串特征,然后会解析PE文件格式,再后调用parsed_features中的函数获取节头特征、导入和导出表特征、文件头特征等,最后合并成一个多维向量,代码如下:
def extract(self, bytez): featurevectors = [fe(bytez) for fe in self.raw_features] binary = lief.PE.parse(bytez) for fe in self.parsed_features: featurevectors.append(fe(binary)) return np.concatenate(featurevectors)
Interface模块基于GBDT模型针对PE文件进行检测。Gym-Malware使用GBDT,在5万个黑样本和5万个白样本上学习,生成的模型保存为文件gradient_boosting.pkl。Interface直接加载该文件,然后对PE文件进行检测。关于GBDT的使用以及恶意程序检测的方法请参考本书第8章“恶意程序检测”相关内容。
Interface模块首先创建全局的PEFeatureExtractor对象feature_extractor,然后加载训练好的GBDT模型文件,创建GBDT分类器对象local_model。定义全局阈值local_model_threshold,当local_model判别的时候,大于该阈值的判定为1,反之判定为0。这里需要指出,该阈值可以根据实际情况调整,通常在0.5~0.9,代码如下:
feature_extractor = PEFeatureExtractor() local_model = joblib.load(os.path.join(module_path, 'gradient_boosting.pkl') ) local_model_threshold = 0.50
Interface模块通过函数get_score_local预测指定字节数组对应的PE文件的分数:
def get_score_local(bytez): features = feature_extractor.extract( bytez ) score = local_model.predict_proba( features.reshape(1,-1) )[0,-1] return score
Interface模块对外提供接口函数get_label_local,用于判定指定字节数组对应的PE文件是否是恶意程序,当分数大于阈值判定为1,反之判定为0,代码如下:
def get_label_local(bytez): score = get_score_local(bytez) label = float( get_score_local(bytez) >= local_model_threshold ) print("score={} (hidden), label={}".format(score,label)) return label
MalwareManipulator模块封装了对PE文件的各种免杀操作,关于免杀技术的介绍请参考本书第9章。
MalwareManipulator模块定义了一个免杀操作的映射表格,用于完成操作名称到具体函数的映射过程,其中包括文件末尾追加随机内容、追加导入表、改变节名称和追加节等,代码如下:
ACTION_TABLE = { 'overlay_append': 'overlay_append', 'imports_append': 'imports_append', 'section_rename': 'section_rename', 'section_add': 'section_add', 'section_append': 'section_append', 'create_new_entry': 'create_new_entry', 'remove_signature': 'remove_signature', 'remove_debug': 'remove_debug', 'upx_pack': 'upx_pack', 'upx_unpack': 'upx_unpack', 'break_optional_header_checksum': 'break_optional_header_checksum', }
具体的转换依靠转换表ACTION_LOOKUP:
ACTION_LOOKUP = {i: act for i, act in enumerate( manipulate.ACTION_TABLE.keys())}
DQNAgent具体实现了强化学习算法,关于DQN的详细介绍请参考本书第6章,本章主要介绍在Gym-Malware中如何使用DQNAgent。
首先定义创建深度学习网络的函数,input_shape指的是输入的特征向量的维度,layers指的是深度学习网络的各层层数,nb_actions指的是动作空间的大小,由于在本例中动作空间是有限的离散值,所以nb_actions事实上也就是动作的个数,同时也是深度学习网络输出层节点数。这里深度学习网络使用的是多层感知机(MLP),所以直接指定层数即可,本例中建议使用两层隐藏层,结点数分别为1024和256,代码如下:
def generate_dense_model(input_shape, layers, nb_actions): model = Sequential() model.add(Flatten(input_shape=input_shape)) model.add(Dropout(0.1)) for layer in layers: model.add(Dense(layer)) model.add(BatchNormalization()) model.add(ELU(alpha=1.0)) model.add(Dense(nb_actions)) model.add(Activation('linear')) return model
然后我们初始化Gym环境,获取环境env的动作空间大小:
ENV_NAME = 'malware-score-v0' env = gym.make(ENV_NAME) nb_actions = env.action_space.n window_length = 1
创建DQNAgent的深度学习网络:
model = generate_dense_model((window_length,) + env.observation_space.shape, layers, nb_actions)
创建策略对象policy,这里使用的是玻尔兹曼算法:
policy = BoltzmannQPolicy()
创建记忆体,大小为32:
memory = SequentialMemory(limit=32, ignore_episode_boundaries=False, window_length=window_length)
创建DQNAgent对象agent,指定使用的深度学习网络、动作空间大小、记忆体、使用的策略和批处理大小等参数,代码如下:
agent = DQNAgent(model=model, nb_actions=nb_actions, memory=memory, nb_steps_warmup=16,enable_double_dqn=True,enable_dueling_network=True, dueling_type='avg',arget_model_update=1e-2, policy=policy, batch_size=16)
编译agent中的深度学习网络并开始学习,学习的总步数为rounds:
agent.compile(RMSprop(lr=1e-3), metrics=['mae']) agent.fit(env, nb_steps=rounds, visualize=False, verbose=2)
MalwareEnv是基于OpenAI Gym开发的,在整个Gym-Malware中是非常重要的角色。如图10-4所示,MalwareEnv对外主要提供了step和reset两个接口,另外还定义了一些内部函数,下面我们重点介绍其功能和实现。
1.Init函数
Init函数主要负责创建MalwareEnv时完成一系列初始化工作,比如初始化动作空间和特征提取的对象PEFeatureExtractor:
self.action_space = spaces.Discrete(len(ACTION_LOOKUP)) self.feature_extractor = pefeatures.PEFeatureExtractor()
图10-4 MalwareEnv主要函数的功能简图
其中,样本的列表保存在变量available_sha256中:
self.available_sha256
完成初始化操作后,需要重置环境:
self._reset()
2.Step函数
Step函数完成了最重要的动作执行、病毒检测以及反馈状态的功能。Step函数的输入通常是数字型,表示的是动作对应的序号action_index。动作的序号可以通过转换表对应成函数名,然后使用函数名调用对应的函数即可。转换表ACTION_LOOKUP是个字典结构,键值为动作对应的序号,内容为对应的函数名称,需要特别注意的是这里的函数名称是个字符串,代码如下:
ACTION_LOOKUP = {i: act for i, act in enumerate( manipulate.ACTION_TABLE.keys())}
当获得了函数名称后,可以通过对应的类的__getattribute__函数获取对应的函数入口,完整的流程如图10-5所示,上述过程封装成内部函数take_action:
_action = MalwareManipulator(bytez).__getattribute__(_action)
当获得了动作对应的函数入口后,执行动作:
图10-5 通过动作序号执行函数的流程图
self._take_action(action_index)
label_function封装杀毒检测的过程,使用其对当前的样本进行检测,获得对应的标记:
self.label = label_function(self.bytez)
如果标记为0,表明已经免杀,设置奖励值reward为10。反之,如果没有超过尝试的步数的阈值,继续循环学习,如果超过了,设置本次学习周期结束。整个Step函数的流程图如图10-6所示,代码如下:
self.observation_space = self.feature_extractor.extract(self.bytez) if self.label == 0: reward = 10.0 # !! a strong reward episode_over = True elif self.turns >= self.maxturns: reward = 0.0 episode_over = True else: reward = 0.0 episode_over = False
3.Reset函数
Reset函数负责重置环境的状态,并随机从样本中选择一个,转换成特征向量后作为初始状态,代码如下:
self.sha256 = random.choice(self.available_sha256) self.observation_space = self.feature_extractor.extract(self.bytez)
图10-6 Step函数的流程图