PentestFusion

Một hệ thống Adapter Fusion cho các tác vụ Pentest, kết hợp nhiều LoRA adapters được fine-tune trên các domain khác nhau của cybersecurity.

📋 Tổng quan

PentestFusion sử dụng kỹ thuật Adapter Fusion để kết hợp 4 LoRA adapters chuyên biệt:

CVE Adapter: Phân tích và khai thác CVE
CAPEC Adapter: Hiểu biết về các attack patterns (CAPEC)
Tools Adapter: Kiến thức về các công cụ pentest
VulHub Adapter: Writeups và hướng dẫn từ VulHub

Mỗi adapter được fine-tune độc lập trên Llama-3.1-8B-Instruct, sau đó được kết hợp thông qua một Fusion Layer sử dụng attention mechanism để tự động điều chỉnh trọng số dựa trên input.

🏗️ Kiến trúc

Base Model

Model: Meta-Llama-3.1-8B-Instruct (4-bit quantization)
Framework: Unsloth (tối ưu hóa training và inference)

LoRA Adapters

Mỗi adapter sử dụng:

Rank (r): 32
LoRA Alpha: 64
Target Modules: q_proj, k_proj, v_proj, o_proj, gate_proj, up_proj, down_proj

Adapter Fusion Layer

Fusion layer sử dụng attention mechanism để:

Tính toán attention scores giữa base model output và từng adapter output
Tự động điều chỉnh trọng số cho từng adapter dựa trên context
Kết hợp các outputs thành một representation thống nhất

# Kiến trúc Fusion
h (base model) → Query
z₁, z₂, z₃, z₄ (adapters) → Keys & Values
Attention(Q, K, V) → Fused Output

📦 Cài đặt

Yêu cầu

Python 3.8+
CUDA-capable GPU (khuyến nghị)
16GB+ RAM

Cài đặt dependencies

pip install -r requirements.txt

🚀 Sử dụng

1. Load model và adapters

from run_fusion import FusionModel
import torch
from unsloth import FastLanguageModel
from peft import PeftModel

# Load base model
BASE_MODEL = "unsloth/Meta-Llama-3.1-8B-Instruct-bnb-4bit"
model, tokenizer = FastLanguageModel.from_pretrained(
    model_name=BASE_MODEL,
    max_seq_length=2048,
)

# Load adapters
model = PeftModel.from_pretrained(model, "./adapters/llama3_cve_adapter", adapter_name="cve")
model.load_adapter("./adapters/llama3_capec_adapter", adapter_name="capec")
model.load_adapter("./adapters/llama3_tool_adapter", adapter_name="tools")
model.load_adapter("./adapters/llama3_vulhub_writeup_adapter", adapter_name="vulhub")

# Tạo fusion model
fusion_model = FusionModel(model, ["cve", "capec", "tools", "vulhub"]).cuda()

# Load fusion layer weights
fusion_model.fusion.load_state_dict(
    torch.load("./fusion_layer.pt", map_location="cuda")
)
fusion_model.eval()

2. Inference

prompt = "Analyze CVE-2017-15715 and outline exploitation steps."

inputs = tokenizer(prompt, return_tensors="pt").to("cuda")

with torch.no_grad():
    logits, weights = fusion_model(**inputs)

# Decode output
output_ids = torch.argmax(logits, dim=-1)
output_text = tokenizer.decode(output_ids[0], skip_special_tokens=True)

# Xem fusion weights (độ quan trọng của từng adapter)
print("Fusion Weights:")
print(f"CVE   : {weights.mean(dim=1)[0][0].item():.4f}")
print(f"CAPEC : {weights.mean(dim=1)[0][1].item():.4f}")
print(f"TOOLS : {weights.mean(dim=1)[0][2].item():.4f}")
print(f"VULHUB: {weights.mean(dim=1)[0][3].item():.4f}")

3. Chạy demo

python run_fusion.py

📁 Cấu trúc thư mục

pentestfusion/
├── adapters/
│   ├── llama3_cve_adapter/          # CVE analysis adapter
│   ├── llama3_capec_adapter/         # CAPEC patterns adapter
│   ├── llama3_tool_adapter/          # Pentest tools adapter
│   └── llama3_vulhub_writeup_adapter/ # VulHub writeups adapter
├── fusion_layer.pt                   # Trained fusion layer weights
├── run_fusion.py                     # Main script
├── requirements.txt                  # Dependencies
└── README.md                         # Documentation

🔬 Chi tiết kỹ thuật

Adapter Fusion Mechanism

Fusion layer sử dụng multi-head attention để kết hợp outputs:

Base model output (h) được dùng làm query
Adapter outputs (z₁...z₄) được dùng làm keys và values
Attention scores được tính toán và normalize bằng softmax
Weighted sum của adapter outputs tạo ra fused representation

Training

Mỗi adapter được fine-tune độc lập trên dataset chuyên biệt
Fusion layer được train sau khi các adapters đã được fine-tune
Training sử dụng mixed precision (FP16/BF16) để tối ưu memory

📊 Adapters

Adapter	Domain	Base Model	Rank	Alpha
CVE	CVE Analysis	Llama-3.1-8B	32	64
CAPEC	Attack Patterns	Llama-3.1-8B	32	64
Tools	Pentest Tools	Llama-3.1-8B	32	64
VulHub	Writeups	Llama-3.1-8B	32	64

🤝 Đóng góp

Contributions are welcome! Please feel free to submit a Pull Request.

📝 License

[Specify your license here]

🙏 Acknowledgments

Unsloth - Fast LLM fine-tuning framework
PEFT - Parameter-Efficient Fine-Tuning
AdapterFusion - Original Adapter Fusion paper
LoRA - Low-Rank Adaptation technique

📚 References

⚠️ Lưu ý

File fusion_layer.pt có kích thước lớn (~192MB). Cân nhắc sử dụng Git LFS hoặc GitHub Releases.
Model yêu cầu GPU với CUDA để inference hiệu quả.
Đảm bảo có đủ VRAM (khuyến nghị 16GB+) để load model và adapters.