CTFer学ML系列

前言

最近闲着学了下深度学习,读了兜哥的数,虽然学ml不能看这本书,但是里边的各种应用例子还是蛮不错的,特别适合上手练习。书中使用的是tensorflow,我是学pytorch入门的,正好闲着把书中的东西用pytorch重写复现一下好了。

这个系列涉及webshell检测、xss攻击向量检测等,边学边练慢慢来。

测了这几个例子,感觉重要的还是数据和特征提取,后边套模型其实都是一把梭。而特征提取就需要对要学习的对象有一定的了解。比如php,兜哥想到的是提取opcode,这样就能形成特征更明显的序列模型,输入RNN的效果一定比直接输入原生的php code效果要好。xss、sqli等等也一样,需要找到更好的特征。

1. RNN检测webshell

模型

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
class BiRNN(nn.Module):
def __init__(self, embedding_dim, num_hiddens, num_layers):
super(BiRNN, self).__init__()
self.embedding = nn.Embedding(token_size, embedding_dim)

# bidirectional为True是双向RNN
self.encoder = nn.LSTM(input_size=embedding_dim,
hidden_size=num_hiddens,
num_layers=num_layers,
bidirectional=True)
# 初始时间步和最终时间步的隐藏状态作为全连接层输入,因此是4*hidden
# 二分类,最后一维就是2,n分类就是n
self.decoder = nn.Linear(4 * num_hiddens, 2)

def forward(self, inputs):
# inputs的形状是(batch_size, token_len),因为LSTM需要将token长度(token_len)作为第一维,
# 所以将输入转置后再提取词特征,输出形状为(seq_len, batch_size, embedding_dim)
embeddings = self.embedding(inputs.permute(1, 0))
# rnn.LSTM只输入embeddings,因此只返回最后一层的隐藏层在各时间步的隐藏状态。
# outputs形状是(token_len, batch_size, 2 * num_hiddens)
outputs, _ = self.encoder(embeddings)
# 连结初始时间步和最终时间步的隐藏状态作为全连接层输入。它的形状为(batch_size, 4 * num_hiddens)
encoding = torch.cat((outputs[0], outputs[-1]), -1)
outs = self.decoder(encoding)
return outs

先贴一下模型代码,这个模型我直接抄了网上情感分类的,然后改了一下输入,整体的处理基本没变,本来想最后fc之前加个softmax后来发现没必要。

就是一个标准的双向LSTM,先将输入用word embedding进行处理,然后是双向LSTM,再一个full connection,不太一样的是torch.cat连结了初始和最终时间步。其实没太懂这里为啥要连结,感觉webshell检测应该和初始时间步没啥关系。

数据处理&特征提取

特征提取方式

对于webshell检测,这应该是最重要的步骤了。
目前常见的输入有两种方式,第一种是直接使用php source输入,第二种是提取php opcode输入。这其实稍微想一下就会感觉第一种没有第二种靠谱,因为输入需要映射到矩阵,而如果输入php源码,那不论怎么映射,感觉特征上都没有映射opcode来的明显,比如call_user_func,参数传入过程中的变换、回溯等等其实根据opcode都能训练出,然而如果根据php source,变量名变一下感觉就够呛了。

因此这里采用opcode的方式,具体提取代码见github。

数据处理

提取opcode后,需要划分测试集、训练集等,这里使用train_test_split划分后,处理成pytorch的DataLoader迭代器,然后就可以根据batch_size进行一轮轮的数据获取并输入训练。

训练

训练的loss采用交叉熵:CrossEntropyLoss,然后反向传播等等就都一样了。
loss使用train_loss_sum / batch_count来计算,train_loss_sum就是每一轮batch_size个训练数据中和预期不符合的数据数量。
准确度的评估见代码,调一下都很清晰了。

2. RNN检测xss攻击向量

Proudly powered by Hexo and Theme by Hacker
© 2021 LFY