redauzhang/common-injection-payload-classfication

Web Attack Detection Model

A CodeBERT-based deep learning model for detecting malicious web requests and payloads. This model can identify SQL injection, XSS, path traversal, command injection, and other common web attack patterns.

Model Description

This model is fine-tuned from microsoft/codebert-base for binary classification of web requests as either benign or malicious.

Model Architecture

• Base Model: CodeBERT (RoBERTa-base architecture)
• Task: Binary Text Classification
• Parameters: 124.6M
• Max Sequence Length: 256 tokens

Performance Metrics

Metric	Training Set	Test Set (125K)	2000-Sample Test
Accuracy	99.30%	99.38%	99.60%
Precision	-	99.47%	99.80%
Recall	-	99.21%	99.40%
F1 Score	-	99.34%	99.60%

Confusion Matrix (Test Set)

	Predicted Benign	Predicted Malicious
Actual Benign	65,914	312
Actual Malicious	464	58,491

Training Details

Dataset

• Total Samples: 625,904
• Training Samples: 500,722 (80%)
• Test Samples: 125,181 (20%)
• Class Distribution: Balanced (47% malicious, 53% benign)
• Sampling Strategy: Balanced sampling with WeightedRandomSampler

Training Configuration

Parameter	Value
Epochs	3
Batch Size	8
Gradient Accumulation Steps	4
Effective Batch Size	32
Learning Rate	2e-5
Warmup Steps	500
Weight Decay	0.01
Max Sequence Length	256
Optimizer	AdamW

Training Progress

Epoch	Train Loss	Train Acc	Test Loss	Test Acc	F1 Score
1	0.0289	98.84%	0.0192	99.09%	0.9904
2	0.0201	99.24%	0.0169	99.08%	0.9903
3	0.0175	99.30%	0.0274	99.38%	0.9934

Hardware

• GPU: NVIDIA Tesla T4 (16GB)
• Training Time: ~24 hours

Model Files

File	Size	Description
best_model.pt	1.4 GB	PyTorch checkpoint (full precision)
model.onnx	476 MB	ONNX model (full precision)
model_quantized.onnx	120 MB	ONNX model (INT8 quantized)

Usage

Quick Start with ONNX Runtime

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

# Load tokenizer and model
tokenizer = RobertaTokenizer.from_pretrained("microsoft/codebert-base")
session = ort.InferenceSession("model_quantized.onnx", providers=['CUDAExecutionProvider', 'CPUExecutionProvider'])

# Predict
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])
        }
    }

# Example usage
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}}

Using 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

# Load model
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)

# Load tokenizer
tokenizer = RobertaTokenizer.from_pretrained("microsoft/codebert-base")

# Predict
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()
    }

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

FastAPI Server

Installation

pip install onnxruntime-gpu transformers fastapi uvicorn pydantic numpy

Start Server

# GPU mode (recommended)
python server_onnx.py --device gpu --quantized --port 8000

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

API Endpoints

Health Check

curl http://localhost:8000/health

Single Prediction

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

Response:

{
  "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
}

Batch Prediction

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 Deployment

GPU Version

# 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 Version

# 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]

Attack Types Detected

This model can detect various web attack patterns including:

Attack Type	Example
SQL Injection	' OR '1'='1' --
Cross-Site Scripting (XSS)	<script>alert(document.cookie)</script>
Path Traversal	../../etc/passwd
Command Injection	; cat /etc/passwd
LDAP Injection	`)(uid=))(
XML Injection	<?xml version="1.0"?><!DOCTYPE foo>
Server-Side Template Injection	{{7*7}}

Limitations

• The model is trained on specific attack patterns and may not detect novel or obfuscated attacks
• Maximum input length is 256 tokens; longer payloads will be truncated
• The model may have false positives on legitimate requests that resemble attack patterns
• Performance may vary on different types of web applications

Ethical Considerations

This model is intended for defensive security purposes only, including:

• Web Application Firewalls (WAF)
• Intrusion Detection Systems (IDS)
• Security monitoring and alerting
• Penetration testing and security assessments

Do not use this model for malicious purposes.

License

This model is released under the MIT License.

Citation

If you use this model in your research or application, please cite:

@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}
}

Acknowledgments

• Microsoft CodeBERT for the pre-trained model
• Hugging Face Transformers for the model framework
• ONNX Runtime for efficient inference