机器学习检测恶意url

尝试下机器学习在安全方面的应用，因为数据集不好找，用的是别人收集的xss数据，实现一个简单的xss攻击检测。
首先处理数据，分别读取正常url和含有xss攻击payload的url，去重并解码，统计数量：

normal = pd.read_csv('normal_examples.csv',header = None)
xssed = pd.read_csv('xssed.csv',header = None)
xssed = list(set(xssed[0]))
normal = list(set(normal[0]))
xssed = [urllib.request.unquote(d) for d in xssed]
normal = [urllib.request.unquote(d) for d in normal]

攻击样本有28775条，正常样本有162637条,数据不平衡，还要注意一下。
之后用正则表达式对数据泛化处理，把数字都替换成0，url中的host都改成u。这步的目的是为了减少分词之后特征的维度。之后用nltk分词，分词的正则大概就是把每个参数名和参数值都切开。最后用空格再把分好的词连起来，方便下一步tfidf计算特征向量。

def word_split(payload,r):
    payload,num=re.subn(r'\d+',"0",payload)
    reurl = r'(?<=://).+?(?=/)'
    payload,num=re.subn(reurl,"u",payload)
    return nltk.regexp_tokenize(payload,r)

r = '(?x)[\w\.]+?\(|\)|"\w+?"|\'\w+?\'|http://\w|</\w+>|<\w+>|<\w+|\w+=|>|[\w\.]+'
all_data = [' '.join(word_split(p,r)) for p in all_data]

因为url也没有什么上下文，就用基于词频的tfidf算法就很适合。直接调用TfidfVectorizer方法生成tfidf特征向量，看了下有160416维。之后再生成标签，用来训练的数据就处理好了。

tfidf_vect = TfidfVectorizer(min_df = 0.0, analyzer="char", sublinear_tf=True, decode_error='ignore',ngram_range=(1,3)) #converting data to vectors
X = tfidf_vect.fit_transform(all_data)
y = [0]*len(normal) + [1] * len(xssed)

然后想找一个baseline做对比，想到很多waf就是基于正则来检测xss，上google找了个召回率比较高的正则跑了一下

def redetect(data):
    xssre = re.compile(r"<(|\/|[^\/>][^>]+|\/[^>][^>]+)>")
    count = 0
    for x in data:
        if xssre.search(x,re.I):
            count += 1
    return count
tp = redetect(xssed)
fp = redetect(normal)
validCount = len(normal)
badCount = len(xssed)
tn = validCount - fp
accuracy = (tp + tn) / (validCount + badCount)
precision = tp / (tp + fp)
recall = tp / badCount
print("Bad samples: %d" % badCount)
print("Good samples: %d" % validCount)
print("Baseline Constant negative: %.6f" % (validCount / (validCount + badCount)))
print("------------")
print("Accuracy: %f" % accuracy)
print("Precision: %f" % precision)
print("Recall: %f" % recall)
print("F1-Score: %f" % (2 * precision * recall / (precision + recall)))

得到结果

Bad samples: 28775
Good samples: 162637
Baseline Constant negative: 0.849670
------------
Accuracy: 0.966058
Precision: 0.971792
Recall: 0.797359
F1-Score: 0.875976

总听说用正则不靠谱，结果确实不咋样，召回率只有八成。
接着还是用机器学习算法跑一下吧，二分类考虑逻辑回归和svm，用lr直接导入tfidf特征，因为数据不平衡，加上class_weight参数：

X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.3, random_state=666)
lgs = LogisticRegression(class_weight="balanced") 
lgs.fit(X_train, y_train)
predicted = lgs.predict(X_test)
fpr, tpr, _ = metrics.roc_curve(y_test, (lgs.predict_proba(X_test)[:, 1]))
auc = metrics.auc(fpr, tpr)

跑完之后再对比下评价指标

Accuracy: 0.998677
Precision: 0.996545
Recall: 0.993683
F1-Score: 0.995112
AUC: 0.999702

确实每项指标都高的不是一星半点，感觉没有跑其他模型的必要了。只是没办法测试是否在真实数据上有过拟合。
想用svm测一下，但是训练不出来，直接报内存错误了，用word2vec来embedding之后还是跑不动，看来svm真的开销比较大。
不过模型总会有漏判误判，对于xss，还是要进行转义、白名单过滤之类的方法。