恶意URL检测的方法很多,这里介绍通过机器学习分析URL文本分词词频来检测恶意URL。训练的数据集为开源数据集,通过机器学习训练检测模型,然后做了部分工程化的应用,将模型持久化,在应用的时候加载进来直接应用,不用重新进行训练。通过接口调用实现恶意URL检测预测判断。
恶意URL检测,对应与机器学习是个分类问题,这里分别用逻辑回归和SVM支持向量机分类模型进行模型实现。
具体实现过程包括 数据载入 –> 数据处理(分词、向量化处理) –> 模型训练 –> 模型保存 –> 模型应用
项目组织结构如下:
一、数据载入
从数据集中载入数据,读取数据,将URL和标签进行识别和区分。
#从文件中获取数据集 def getDataFromFile(filename='data/data.csv'): input_url = filename data_csv = pd.read_csv(input_url, ',', error_bad_lines=False) data_df = pd.DataFrame(data_csv) url_df = np.array(data_df) random.shuffle(url_df) y = [d[1] for d in url_df] inputurls = [d[0] for d in url_df] return inputurls,y
二、数据处理(分词、向量化处理)
数据处理实现对URL的分词及向量化处理
分词:分析URL根据,.-进行分词,由于com、cn等常用域名不是关键影响因素,所以分词的时候去掉了
分词
#分词 def getTokens(input): web_url = input.lower() urltoken = [] dot_slash = [] slash = str(web_url).split('/') for i in slash: r1 = str(i).split('-') token_slash = [] for j in range(0, len(r1)): r2 = str(r1[j]).split('.') token_slash = token_slash + r2 dot_slash = dot_slash + r1 + token_slash urltoken = list(set(dot_slash)) if 'com' in urltoken: urltoken.remove('com') if 'cn' in urltoken: urltoken.remove('cn') return urltoken
向量化处理
将分词以后的结果进行词频的向量化处理,形成可以用于模型训练的稀疏矩阵向量
all_urls,y=getDataFromFile(datapath) url_vectorizer = TfidfVectorizer(tokenizer=getTokens) x = url_vectorizer.fit_transform(all_urls)
三、模型训练
将经过处理后的训练数据用模型进行训练,将数据集分为两部分一部分用于训练,一部分用于测试评估。
#训练,通过逻辑回归模型训练 def trainLR(datapath): all_urls,y=getDataFromFile(datapath) url_vectorizer = TfidfVectorizer(tokenizer=getTokens) x = url_vectorizer.fit_transform(all_urls) x_train, x_test, y_train, y_test = train_test_split(x, y, test_size=0.2, random_state=42) l_regress = LogisticRegression() # Logistic regression l_regress.fit(x_train, y_train) l_score = l_regress.score(x_test, y_test) print("score: {0:.2f} %".format(100 * l_score)) return l_regress,url_vectorizer
用逻辑回归模型训练的结果是 score: 98.50 %
#训练,通过SVM支持向量机模型训练 def trainSVM(datapath): all_urls, y = getDataFromFile(datapath) url_vectorizer = TfidfVectorizer(tokenizer=getTokens) x = url_vectorizer.fit_transform(all_urls) x_train, x_test, y_train, y_test = train_test_split(x, y, test_size=0.2, random_state=42) svmModel=svm.LinearSVC() svmModel.fit(x_train, y_train) svm_score=svmModel.score(x_test, y_test) print("score: {0:.2f} %".format(100 * svm_score)) return svmModel,url_vectorizer
用SVM模型训练的结果是 score: 99.64 %
可以看出SVM模型训练的结果比逻辑回归模型训练的效果要稍好。
四、保存模型
将训练好的模型进行持久化保存,通过pickle.dump()的方式把训练好的模型参数及特征保存至模型文件,以便于应用的时候不要再进行训练,直接应用训练好的模型。
#保存模型及特征 def saveModel(model,vector): #保存模型 file1 = modelfile_path with open(file1, 'wb') as f: pickle.dump(model, f) f.close() #保存特征 file2 = vectorfile_path with open(file2, 'wb') as f2: pickle.dump(vector, f2) f2.close()
通过main方法执行训练模型及保存模型
if __name__ == '__main__': #model,vector=trainLR('data/data.csv') model, vector = trainSVM('data/data.csv') saveModel(model,vector)
四、模型应用
通过pickle.load载入已经训练好的模型和特征,并用Flask暴露一个接口调用模型的预测方法进行预测。
载入已经训练好的模型
#载入已经训练好的模型 def loadModel(): file1 = modelfile_path with open(file1, 'rb') as f1: model = pickle.load(f1) f1.close() file2 = vectorfile_path with open(file2, 'rb') as f2: vector = pickle.load(f2) f2.close() return model,vector
通过接口进行调用
#通过接口进行调用 @app.route('/<path:path>') def show_predict(path): X_predict = [] X_predict.append(path) model, vector = loadModel() x = vector.transform(X_predict) y_predict = model.predict(x) print(y_predict[0]) return "url predict: "+str(y_predict[0])
Be First to Comment