机器学习实现恶意URL检测实战一

恶意URL检测的方法很多，这里介绍通过机器学习分析URL文本分词词频来检测恶意URL。训练的数据集为开源数据集，通过机器学习训练检测模型，然后做了部分工程化的应用，将模型持久化，在应用的时候加载进来直接应用，不用重新进行训练。通过接口调用实现恶意URL检测预测判断。

恶意URL检测，对应与机器学习是个分类问题，这里分别用逻辑回归和SVM支持向量机分类模型进行模型实现。

具体实现过程包括数据载入–>数据处理（分词、向量化处理）–>模型训练–>模型保存–>模型应用

项目组织结构如下：

一、数据载入

从数据集中载入数据，读取数据，将URL和标签进行识别和区分。

#从文件中获取数据集defgetDataFromFile(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等常用域名不是关键影响因素，所以分词的时候去掉了

分词

#分词defgetTokens(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 inrange(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)

三、模型训练

将经过处理后的训练数据用模型进行训练，将数据集分为两部分一部分用于训练，一部分用于测试评估。

#训练,通过逻辑回归模型训练deftrainLR(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支持向量机模型训练deftrainSVM(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()的方式把训练好的模型参数及特征保存至模型文件，以便于应用的时候不要再进行训练，直接应用训练好的模型。

#保存模型及特征defsaveModel(model,vector):#保存模型
    file1 = modelfile_path
    withopen(file1,'wb')as f:
        pickle.dump(model, f)
    f.close()#保存特征
    file2 = vectorfile_path
    withopen(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暴露一个接口调用模型的预测方法进行预测。

载入已经训练好的模型

#载入已经训练好的模型defloadModel():
    file1 = modelfile_path
    withopen(file1,'rb')as f1:
        model = pickle.load(f1)
    f1.close()

    file2 = vectorfile_path
    withopen(file2,'rb')as f2:
        vector = pickle.load(f2)
    f2.close()return model,vector

通过接口进行调用

#通过接口进行调用@app.route('/<path:path>')defshow_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])