PentestFusion

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 để:

1. Tính toán attention scores giữa base model output và từng adapter output
2. Tự động điều chỉnh trọng số cho từng adapter dựa trên context
3. 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:

1. Base model output (h) được dùng làm query
2. Adapter outputs (z₁...z₄) được dùng làm keys và values
3. Attention scores được tính toán và normalize bằng softmax
4. 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

🤝 Đó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

• AdapterFusion: Non-Destructive Task Composition for Transfer Learning
• LoRA: Low-Rank Adaptation of Large Language Models
• Bayesian-LoRA: Optimal Quantization and Rank Values

⚠️ 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.