首页 > Posts > 回顾项目-用机器学习检测恶意流量

回顾项目-用机器学习检测恶意流量

发布时间 2025-07-15 14:59:13 / 次阅读

翻出来之前在一个安全中心做的恶意流量识别项目的代码, 做下总结

使用 PySpark 和 MLlib 构建 Web 攻击检测模型：从 N-gram 到逻辑回归

在网络安全领域，检测恶意 Web 请求（如 SQL 注入、跨站脚本攻击 XSS）是一项至关重要的任务。本文将深入探讨如何使用 PySpark 和 MLlib 构建一个高效的 Web 攻击检测模型。我们将从数据预处理、特征工程（N-gram 分词）、模型训练到最终的预测和评估，一步步揭示其背后的技术细节。

技术栈概览

我们的解决方案基于以下核心技术：

  • PySpark: 用于大规模数据处理和分布式计算。
  • MLlib: PySpark 的机器学习库，提供丰富的算法和工具，包括特征转换、模型训练和评估。
  • N-gram 分词: 一种文本特征提取技术，将 URL 和 Payload 切分成连续的子字符串，用于捕捉恶意模式。
  • 逻辑回归 (Logistic Regression): 一种常用的分类算法，用于预测 Web 请求是否为恶意。

数据准备与特征工程

首先，我们需要准备用于训练和预测的数据。通常，这些数据来源于 Web 服务器的访问日志，包含 URL、POST 数据、HTTP 请求头等信息。

1. 加载数据

我们使用 PySpark 从 Parquet 格式的文件中读取数据：

1data = spark.read.parquet("data/new_data").fillna("").repartition(1024)
2data.cache()

fillna("") 用于处理缺失值，repartition(1000) 用于将数据分成多个分区，以便并行处理。cache() 则将数据缓存到内存中，加速后续操作。

2. 特征工程：N-gram 分词(可以用现有的大模型分词代替, o200k )

特征工程是机器学习的关键步骤。在本例中，我们使用 N-gram 分词技术从 URL 和 POST 数据中提取特征。

1from pyspark.sql.functions import udf
2from pyspark.sql.types import StringType, ArrayType
3import pyspark.sql.functions as F
4
5# 自定义 UDF 函数，用于 N-gram 分词
6udf_ngram_split = udf(ngram_split, ArrayType(StringType()))
7
8sel_data = data.select("url","post",udf_ngram_split("url","post").alias("data"),F.lit(1).alias("label"))

ngram_split 是一个自定义的 Python 函数，它将输入的字符串（URL 和 POST 数据）切分成 N-gram 序列。例如，对于字符串 “path”，3-gram 分词的结果是 ["pat", "ath"]。

3. 其他特征

除了 N-gram 分词，我们还可以提取其他有用的特征，例如：

  • HTTP 请求方法 (GET, POST, HEAD)
  • HTTP 响应状态码 (2xx, 3xx, 4xx, 5xx)
  • URL 的信息熵
  • User-Agent 是否为已知恶意 UA
  • 请求路径是否为常见的扫描器路径

这些特征可以通过自定义 UDF 函数提取，并添加到 DataFrame 中。

模型训练与预测

在完成特征工程后，我们可以使用 MLlib 训练一个逻辑回归模型。

1. 构建 Pipeline

MLlib 的 Pipeline 提供了一种方便的方式来组织机器学习流程。一个典型的 Pipeline 包含以下步骤：

  • Tokenizer: 将文本数据切分成单词或 N-gram。
  • HashingTF 或 CountVectorizer: 将文本特征转换为数值型向量。
  • IDF (可选): 计算逆文档频率，进行特征加权。
  • Logistic Regression: 训练逻辑回归模型。

2. 加载预训练模型

由于篇幅限制，本文不详细介绍模型训练过程。我们假设已经训练好了一个模型，并将其保存到磁盘。

1from pyspark.ml import PipelineModel
2
3model = PipelineModel.load("cv_lr_model")

3. 模型预测

加载模型后，我们可以使用它来预测新的 Web 请求是否为恶意。

1test = model.transform(sel_data)
2test.groupby("prediction").count().show()

transform 方法将 Pipeline 应用于输入数据，生成包含预测结果的 DataFrame。groupby("prediction").count().show() 用于统计预测结果，例如：

+----------+------+
|prediction| count|
+----------+------+
|       0.0| 11580|
|       1.0|831150|
+----------+------+

4. 模型评估

为了评估模型的性能，我们需要使用一些指标，例如准确率、召回率、F1 值等。

1# 假设所有数据都是正样本
2TP = 831150  # True Positives
3FN = 11580  # False Negatives
4
5recall = TP / (TP + FN)
6print(f"Recall: {recall}") # 召回率

在本例中，我们假设所有输入数据都是已知的恶意样本，因此我们主要关注模型的召回率。召回率越高，说明模型能够识别出更多的恶意请求。

总结与展望

本文介绍了如何使用 PySpark 和 MLlib 构建一个 Web 攻击检测模型。通过 N-gram 分词和逻辑回归，我们可以有效地识别恶意 Web 请求。

未来，我们可以进一步优化模型，例如：

  • 使用更复杂的特征工程技术，如 Word2Vec 或 BERT。
  • 使用更大的数据集进行训练，提高模型的泛化能力。

希望本文能够帮助读者了解如何使用 PySpark 和 MLlib 构建 Web 攻击检测模型，并为网络安全领域的研究和实践提供一些思路。

希望这篇博客文章对你有所帮助！