Web 攻击检测模型

一个基于 CodeBERT 的深度学习模型，用于检测恶意 Web 请求和攻击载荷。该模型可以识别 SQL 注入、XSS、路径遍历、命令注入以及其他常见的 Web 攻击模式。

模型描述

该模型基于 microsoft/codebert-base 进行微调，用于对 Web 请求进行二分类，判断其为良性或恶意。

模型架构

• 基础模型：CodeBERT（RoBERTa-base 架构）
• 任务：二分类文本分类
• 参数量：124.6M
• 最大序列长度：256 tokens

性能指标

混淆矩阵（测试集）

训练详情

数据集

• 总样本数：625,904
• 训练样本数：500,722（80%）
• 测试样本数：125,181（20%）
• 类别分布：平衡（47% 恶意，53% 良性）
• 采样策略：使用 WeightedRandomSampler 进行平衡采样

训练配置

训练进度

硬件

• GPU：NVIDIA Tesla T4（16GB）
• 训练时间：约 24 小时

模型文件

使用方法

使用 ONNX Runtime 快速开始

import numpy as np
import onnxruntime as ort
from transformers import RobertaTokenizer

# 加载分词器和模型
tokenizer = RobertaTokenizer.from_pretrained(\"microsoft/codebert-base\")
session = ort.InferenceSession(\"model_quantized.onnx\", providers=['CUDAExecutionProvider', 'CPUExecutionProvider'])

# 预测
def predict(payload: str) -> dict:
    inputs = tokenizer(
        payload,
        max_length=256,
        padding='max_length',
        truncation=True,
        return_tensors='np'
    )
    
    outputs = session.run(
        None,
        {
            'input_ids': inputs['input_ids'].astype(np.int64),
            'attention_mask': inputs['attention_mask'].astype(np.int64)
        }
    )
    
    probs = outputs[0][0]
    pred_idx = np.argmax(probs)
    
    return {
        \"prediction\": \"malicious\" if pred_idx == 1 else \"benign\",
        \"confidence\": float(probs[pred_idx]),
        \"probabilities\": {
            \"benign\": float(probs[0]),
            \"malicious\": float(probs[1])
        }
    }

# 示例用法
result = predict(\"SELECT * FROM users WHERE id=1 OR 1=1--\")
print(result)
# {'prediction': 'malicious', 'confidence': 0.9355, 'probabilities': {'benign': 0.0645, 'malicious': 0.9355}}

使用 PyTorch

import torch
import torch.nn as nn
from transformers import RobertaTokenizer, RobertaModel

class CodeBERTClassifier(nn.Module):
    def __init__(self, model_path=\"microsoft/codebert-base\", num_labels=2, dropout=0.1):
        super().__init__()
        self.codebert = RobertaModel.from_pretrained(model_path)
        self.dropout = nn.Dropout(dropout)
        self.classifier = nn.Linear(self.codebert.config.hidden_size, num_labels)
    
    def forward(self, input_ids, attention_mask):
        outputs = self.codebert(input_ids=input_ids, attention_mask=attention_mask)
        pooled_output = outputs.pooler_output
        pooled_output = self.dropout(pooled_output)
        logits = self.classifier(pooled_output)
        return logits

# 加载模型
device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")
model = CodeBERTClassifier()
model.load_state_dict(torch.load(\"best_model.pt\", map_location=device))
model.eval()
model.to(device)

# 加载分词器
tokenizer = RobertaTokenizer.from_pretrained(\"microsoft/codebert-base\")

# 预测
def predict(payload: str) -> dict:
    inputs = tokenizer(
        payload,
        max_length=256,
        padding='max_length',
        truncation=True,
        return_tensors='pt'
    ).to(device)
    
    with torch.no_grad():
        logits = model(inputs['input_ids'], inputs['attention_mask'])
        probs = torch.softmax(logits, dim=-1)[0]
    
    pred_idx = torch.argmax(probs).item()
    
    return {
        \"prediction\": \"malicious\" if pred_idx == 1 else \"benign\",
        \"confidence\": probs[pred_idx].item()
    }

# 示例
result = predict(\"<script>alert('xss')</script>\")
print(result)
# {'prediction': 'malicious', 'confidence': 0.9998}

FastAPI 服务器

安装

pip install onnxruntime-gpu transformers fastapi uvicorn pydantic numpy

启动服务器

# GPU 模式（推荐）
python server_onnx.py --device gpu --quantized --port 8000

# CPU 模式
python server_onnx.py --device cpu --quantized --port 8000

API 端点

健康检查

curl http://localhost:8000/health

单条预测

curl -X POST http://localhost:8000/predict \
  -H \"Content-Type: application/json\" \
  -d '{\"payload\": \"SELECT * FROM users WHERE id=1 OR 1=1--\"}'

响应：

{
  \"payload\": \"SELECT * FROM users WHERE id=1 OR 1=1--\",
  \"prediction\": \"malicious\",
  \"confidence\": 0.9355,
  \"probabilities\": {\"benign\": 0.0645, \"malicious\": 0.9355},
  \"inference_time_ms\": 15.23
}

批量预测

curl -X POST http://localhost:8000/batch_predict \
  -H \"Content-Type: application/json\" \
  -d '{\"payloads\": [\"<script>alert(1)</script>\", \"GET /api/users HTTP/1.1\"]}'

Docker 部署

GPU 版本

# Dockerfile
FROM nvidia/cuda:11.8-cudnn8-runtime-ubuntu22.04

RUN apt-get update && apt-get install -y python3 python3-pip
RUN pip3 install onnxruntime-gpu transformers fastapi uvicorn pydantic numpy

WORKDIR /app
COPY model_quantized.onnx ./models/
COPY server_onnx.py .

EXPOSE 8000
CMD [\"python3\", \"server_onnx.py\", \"--device\", \"gpu\", \"--quantized\"]

CPU 版本

# Dockerfile.cpu
FROM python:3.10-slim

RUN pip install onnxruntime transformers fastapi uvicorn pydantic numpy

WORKDIR /app
COPY model_quantized.onnx ./models/
COPY server_onnx.py .

EXPOSE 8000
CMD [\"python\", \"server_onnx.py\", \"--device\", \"cpu\", \"--quantized\"]

Docker Compose

version: '3.8'
services:
  web-attack-detector:
    build: .
    ports:
      - \"8000:8000\"
    deploy:
      resources:
        reservations:
          devices:
            - driver: nvidia
              count: 1
              capabilities: [gpu]

可检测的攻击类型

该模型可以检测多种 Web 攻击模式，包括：

局限性

• 该模型基于特定攻击模式训练，可能无法检测新型或混淆的攻击
• 最大输入长度为 256 个 tokens，更长的载荷将被截断
• 该模型可能会对类似攻击模式的合法请求产生误报
• 在不同类型的 Web 应用上性能可能会有所差异

伦理考量

该模型仅用于防御性安全目的，包括：

• Web 应用防火墙（WAF）
• 入侵检测系统（IDS）
• 安全监控和告警
• 渗透测试和安全评估

请勿将此模型用于恶意目的。

许可证

该模型采用 MIT 许可证发布。

引用

如果您在研究或应用中使用该模型，请引用：

@misc{web-attack-detection-codebert,
  author = {Your Name},
  title = {Web Attack Detection Model based on CodeBERT},
  year = {2024},
  publisher = {Hugging Face},
  howpublished = {\url{https://huggingface.co/your-username/web-attack-detection}},
  note = {Fine-tuned CodeBERT model for detecting malicious web requests}
}

@article{feng2020codebert,
  title = {CodeBERT: A Pre-Trained Model for Programming and Natural Languages},
  author = {Feng, Zhangyin and Guo, Daya and Tang, Duyu and Duan, Nan and Feng, Xiaocheng and Gong, Ming and Shou, Linjun and Qin, Bing and Liu, Ting and Jiang, Daxin and Zhou, Ming},
  journal = {Findings of the Association for Computational Linguistics: EMNLP 2020},
  year = {2020},
  pages = {1536--1547},
  doi = {10.18653/v1/2020.findings-emnlp.139}
}

@article{liu2019roberta,
  title = {RoBERTa: A Robustly Optimized BERT Pretraining Approach},
  author = {Liu, Yinhan and Ott, Myle and Goyal, Naman and Du, Jingfei and Joshi, Mandar and Chen, Danqi and Levy, Omer and Lewis, Mike and Zettlemoyer, Luke and Stoyanov, Veselin},
  journal = {arXiv preprint arXiv:1907.11692},
  year = {2019}
}

致谢

• Microsoft CodeBERT 提供预训练模型
• Hugging Face Transformers 提供模型框架
• ONNX Runtime 提供高效推理