# 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

```python
# 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

```bash
pip install -r requirements.txt
```

## 🚀 Sử dụng

### 1. Load model và adapters

```python
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

```python
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

```bash
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

| 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](https://github.com/unslothai/unsloth) - Fast LLM fine-tuning framework
- [PEFT](https://github.com/huggingface/peft) - Parameter-Efficient Fine-Tuning
- [AdapterFusion](https://arxiv.org/abs/2005.00247) - Original Adapter Fusion paper
- [LoRA](https://arxiv.org/abs/2106.09685) - Low-Rank Adaptation technique

## 📚 References

- [AdapterFusion: Non-Destructive Task Composition for Transfer Learning](https://arxiv.org/abs/2005.00247)
- [LoRA: Low-Rank Adaptation of Large Language Models](https://arxiv.org/abs/2106.09685)
- [Bayesian-LoRA: Optimal Quantization and Rank Values](https://arxiv.org/pdf/2406.13046)

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