源代码

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
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
# 数据提取&处理（csv）
import pandas as pd
from sklearn.utils import shuffle
# 特征提取
import re
from sklearn.feature_extraction.text import TfidfVectorizer
# 训练集和测试集划分
from sklearn.model_selection import train_test_split
# 逻辑回归算法训练模型
from sklearn.linear_model import LogisticRegression
# 模型评估报告打印
from sklearn.metrics import classification_report
# 模型保存
import pickle


def read_dmoz0409(file_path, lable):
    normal_pd = pd.read_csv(file_path, header=None, names=["url"], usecols=[1], nrows=14989)
    normal_pd["lable"] = lable
    # print(normal_pd)

    return normal_pd


def read_phishing_verified_online(file_path, lable):
    malicious_pd = pd.read_csv(file_path, usecols=[1], nrows=14989)
    malicious_pd["lable"] = lable
    # print(malicious_pd)

    return malicious_pd


def data_process(normal, malicious):
    all = pd.concat([normal, malicious])
    data = all["url"]
    lable = all["lable"]
    data, lable = shuffle(data, lable, random_state=42)


    return data, lable


def data_tokenizer(data):
    return re.findall(r'\w+', data)


def data_vectorizer(data, lable):
    vectorizer = TfidfVectorizer(tokenizer=data_tokenizer)
    x = vectorizer.fit_transform(data)
    y = lable

    return x, y, vectorizer


def model_train(x_train, y_train):
    model = LogisticRegression()
    model.fit(x_train, y_train)

    return model


def model_test(model, x_test, y_test):
    # accuracy = model.score(x_test, y_test)
    # print("模型准确率:", accuracy)
    y_predict = model.predict(x_test)
    report = classification_report(y_test, y_predict, labels=["malicious", "normal"], target_names=["恶意URL", "正常URL"],digits=2)

    return report


def model_save(model, vectorizer, report):
    save = input("是否保存训练的模型（y/n）：")
    if save == "y" or save == "Y":
        with open('model.pkl', 'wb') as f:
            pickle.dump(model, f)
        with open('vectorizer.pkl', 'wb') as f:
            pickle.dump(vectorizer, f)
        with open('report.txt', 'w') as f:
            f.write(report)
        print("保存成功！")
        return True
    else:
        return False


def new_model():
    # 数据提取
    print("数据提取...")
    normal = read_dmoz0409("dmoz0409.csv", "normal")
    malicious = read_phishing_verified_online("phishing_verified_online.csv", "malicious")

    # 数据处理（合并，打乱数据集）
    print("数据处理...")
    data, lable = data_process(normal, malicious)

    # 特征提取
    print("向量化...")
    x, y, vectorizer = data_vectorizer(data, lable)
    # print(vectorizer.vocabulary_)
    # print(x.toarray)

    # 数据集划分
    x_train, x_test, y_train, y_test = train_test_split(x, y, test_size=0.3)

    # 模型训练
    print("模型训练...")
    model = model_train(x_train, y_train)

    # 模型评估
    report = model_test(model, x_test, y_test)
    print("模型评估报告：" + "\n", report)

    # 模型保存?
    model_save(model, vectorizer, report)


def use_model():
    while True:
        url_list = input("输入需要测试的url：").split("\n")
        with open("vectorizer.pkl", "rb") as f:
            vectorizer = pickle.load(f)
        x = vectorizer.transform(url_list)
        with open("model.pkl", "rb") as f:
            model = pickle.load(f)
        y_predict = model.predict(x)
        for i in range(len(url_list)):
            print(url_list[i], y_predict[i])

if __name__ == "__main__":
    choice = input("训练新的模型或是使用已经训练的模型（train/use）：")
    if choice == "train":
        new_model()
    elif choice=="use":
        print("-------恶意URL检测-------")
        use_model()
    else:
        print("无效的输入！")