利用决策树检测web攻击

发表于 | 分类于

前言

       AI目前在各个领域的使用都非常广泛，安全领域多年前就想引入AI来帮助加强威胁检测能力，但由于安全问题的多变性和多面性，我目前觉得AI在安全领域的能力还没有得到很好地表现，目前AI本身主要是基于样本学习，安全威胁的样本太过于多样，想要收集这些样本只有大体量公司的才能做到，或者类似威胁情报做法，大家共享样本，不过数据的来源主要还是依托于巨头公司的业务，在目前数据就是资产的环境下， 我想没有哪家巨头愿意共享出自己的数据，其次样本的分析也需要大量安全专家长时间的工作累积，所以感觉AI安全还是有很长的一段路要走。
       刚好最近公司大佬在研究AI算法检测安全威胁，也跟着学习了一些，想自己做个利用决策树检测web攻击的模型，直接上代码。

实战代码

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
#!/usr/bin/env python
# -*- coding: utf-8 -*-
# Author: Lawliet

import os
from sklearn import tree
from sklearn.feature_extraction.text import TfidfVectorizer


X = []
Y = []
with open("./sample/xss-20000.txt", 'r') as f:
    for line in f:
        X.append(line.strip())
        Y.append(0)
with open("./sample/good-xss-200000.txt", 'r') as f:
    for line in f:
        X.append(line.strip())
        Y.append(1)
		
tfidf2 = TfidfVectorizer()
Xr = tfidf2.fit_transform(X)
clf = tree.DecisionTreeClassifier()
clf = clf.fit(Xr, Y)
print "Model build success"

Z = ['<script>alert(1)</script>','a','script','alert','xss','open','onerror=alert(1)']
print "Test sample list:"
print Z
print '\n'
Zr = tfidf2.transform(Z)

print "detection..."
for i in range(len(Z)):
	if clf.predict(Zr)[i] == 0:
		print "alert_payload:  " + Z[i]

运行结果

image
       后面想订阅公司kafuka的队列，通过模型来检测web攻击，并尝试优化模型的检测能力和测试误报率。