【AI】浅析恶意文件静态检测及部分问题解决思路 前言 分析

发表于

本文正在参加「金石计划 . 瓜分6万现金大奖」

前言

随着互联网的繁荣和发展，海量的核心数据和网络应用也不断向云端、数据中心等关键信息基础设施整合和迁移，主机安全也因此成为网络攻防战的焦点。恶意文件 是指由攻击者专门设计的，在未经所有者许可的情况下用来访问计算机、损害或破坏系统，对保密性、完整性或可用性进行攻击的文件，是当前互联网安全的主要威胁之一。目前，比较主流的恶意文件包括恶意脚本、漏洞利用、蠕虫、木马和间谍软件以及他们的组合或变体。

为了应对挑战，恶意文件静态检测的思想被提了出来。基于机器学习算法的防护技术为实现高准确率、自动化的未知恶意文件检测提供了行之有效的技术途径，已逐渐成为业内研究的热点。

接下来博主将简单介绍其中一种恶意文件静态检测模型的部分内容；

番外：对于想了解梯度下降算法的小伙伴，也可以看看博主的往期博文：

分析

这里的恶意文件静态检测是将恶意文件的二进制转成灰度图，作为 CNN 模型的输入，经过一系列的过程得到输出，然后进行对比、评估等；

考虑到每个样本的大小是不固定的，本来是以 1M 大小作为区分界限的，这里的话，使用 PadSequence 来确保数据长度的一致性；

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
py复制代码class PadSequence(object):

    ...
    
    def pltexe(self, arr):
        arr_n = len(arr) // (1024*1024)
        arr_end_len = len(arr) % (1024*1024)
        re_arr = []
        siz = 1024
        
        # 矩阵转换：行列变化，总数不变
        for ite in range(arr_n):
            st = ite * 1024 * 1024
            pggg0 = np.array(arr[st : st+1024*1024])
            re_arr.append(pggg0.reshape(siz,siz) / 255)
        
        # 用 0 补足
        if arr_end_len!= 0 :
            arr_ = (1024*1024-arr_end_len) * [0]
            pggg0 = np.array(arr[1024*1024*arr_n:] + arr_)
            re_arr.append(pggg0.reshape(siz,siz) / 255)
 
        return re_arr
        
    def doooo_(self, filelist):
        ...
        
        # 设定列表长度不超过20
        if len(featurelist) > 20:
            re_feature_lab = random.sample(list(zip(featurelist,labellist)), 20)
            featurelist = [x[0] for x in re_feature_lab]
            labellist = [x[1] for x in re_feature_lab]
        
        ...
        
        return featurelist_batch, labellist_batch
    
    def __call__(self, batch):
        return self.doooo_(batch)

然后进行数据加载：

1
2
3
4
5
6
7
py复制代码with open(path, 'rb') as f:
        train_data = pickle.load(f)

train_loader = DataLoader(train_data, batch_size=10, shuffle=True, num_workers = 20, collate_fn=PadSequence(maxlen = 0))

pad = PadSequence()
pad.__call__(train_data[:4])[0]

image.png

最后进入模型进行训练以及验证；

image.png

TIP

在模型训练中，可以使用 try...excpet 模块，即使因为意外中断训练，之前的训练结果也都保存下来了，下次训练就不用重头开始了：

1
2
3
4
5
6
py复制代码try:
    ...
    
except:
    model = model.eval()
    torch.save(model.state_dict(), 'error.pth')

问题解决

OOM

在启动项目时，可能会出现以下报错：

1
2
3
4
5
6
7
8
sql复制代码(sid10t) bash-4.2# python model_run.py 
Traceback (most recent call last):
  ...
RuntimeError: Caught RuntimeError in replica 0 on device 0.

Original Traceback (most recent call last):
  ...
RuntimeError: CUDA out of memory. Tried to allocate 24.00 MiB (GPU 0; 7.80 GiB total capacity; 6.31 GiB already allocated; 6.56 MiB free; 6.46 GiB reserved in total by PyTorch)

这是因为我们将 DataLoader 里的 batch_size 参数设置的过大了，从而导致了显存溢出；

那么无非就是两个解决方案：

  • 多选定几个 CUDA；
  • 将 batch_size 参数调小；

对于第一种方案，可以一股脑的将机子上的所有 CUDA 全部选上：

1
2
3
py复制代码device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = model.to(device)
model= nn.DataParallel(model)

对于第二种方案，将 batch_size 参数调小，也是有讲究的，我们要尽可能的提高资源的利用率，因此需要做一些操作：

  1. 首先是查看模型占用了多少 GPU，watch -n 1 nvidia-smi：
    image.png
  2. 然后折半减少 batch_size，查看显存占用率，调节至合适大小：
    image.png

Socket exception

由于模型跑在空闲的机子上，而样本却在另一台机子上，因此，需要通过 SFTP 进行读取，不出意外的话，要出意外了；

首先是在 pad 函数里构建 SFTP 连接：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
py复制代码class PadSequence(object):
    def __init__(self, maxlen = 8000):
        self.maxlen = maxlen
        ...
        self.client = paramiko.SSHClient()
        self.client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
        self.client.connect(self.hostname, self.port ,self.username, self.password, compress=True)
        self.sftp_client = self.client.open_sftp()
    
    def getfile_ftp(self, file_path):
        remote_file = self.sftp_client.open(file_path, 'rb')
        try:
            str_object_with_pe_file_data = remote_file.read()
        finally:
            remote_file.close()
        return str_object_with_pe_file_data

    def __call__(self, batch):
        return self.doooo_(batch)

然后在 DataLoader 中使用到它的回显函数：

1
py复制代码train_loader = DataLoader(train_data, batch_size=4, shuffle=True, num_workers = 20, collate_fn=PadSequence(maxlen = 0))

好的，做完上面的之后，坑爹的来了，报错了：Socket exception: Connection reset by peer (104)

1668414245075_3148F4BD-D120-4d3c-8640-BB3DDCEFC492.png

不知道是因为 DataLoader 底层逻辑问题，还是这台服务器的问题，症结就是在于不能使用多进程进行 SFTP 读取，因此这里的解决方案就是将参数 num_workers 置为 0；

后记

以上就是 浅析恶意文件静态检测及部分问题解决思路 的全部内容了，大致讲述了恶意文件静态检测的其中一种思路，以及图文结合的分析了部分问题的解决思路，希望大家有所收获！

📝 上篇精讲：【AI】浅谈梯度下降算法（实战篇）

💖 我是 𝓼𝓲𝓭𝓲𝓸𝓽，期待你的关注；

👍 创作不易，请多多支持；

🔥 系列专栏：AI 项目实战

本文转载自: 掘金

开发者博客 – 和开发相关的 这里全都有

0%