mirror of
https://github.com/CyberSecurityUP/NeuroSploit.git
synced 2026-05-27 16:42:22 +02:00
NeuroSploit v3.2 - Autonomous AI Penetration Testing Platform
116 modules | 100 vuln types | 18 API routes | 18 frontend pages Major features: - VulnEngine: 100 vuln types, 526+ payloads, 12 testers, anti-hallucination prompts - Autonomous Agent: 3-stream auto pentest, multi-session (5 concurrent), pause/resume/stop - CLI Agent: Claude Code / Gemini CLI / Codex CLI inside Kali containers - Validation Pipeline: negative controls, proof of execution, confidence scoring, judge - AI Reasoning: ReACT engine, token budget, endpoint classifier, CVE hunter, deep recon - Multi-Agent: 5 specialists + orchestrator + researcher AI + vuln type agents - RAG System: BM25/TF-IDF/ChromaDB vectorstore, few-shot, reasoning templates - Smart Router: 20 providers (8 CLI OAuth + 12 API), tier failover, token refresh - Kali Sandbox: container-per-scan, 56 tools, VPN support, on-demand install - Full IA Testing: methodology-driven comprehensive pentest sessions - Notifications: Discord, Telegram, WhatsApp/Twilio multi-channel alerts - Frontend: React/TypeScript with 18 pages, real-time WebSocket updates
This commit is contained in:
CyberSecurityUP
@@ -0,0 +1,372 @@
|
||||
"""
|
||||
NeuroSploit v3 - Agent Task Manager
|
||||
|
||||
Sub-task spawning and tracking system with priority queue.
|
||||
Enables concurrent task execution with dependency awareness.
|
||||
|
||||
Inspired by CAI framework's agent-as-tool and task delegation patterns.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import uuid
|
||||
import time
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Dict, List, Any, Optional, Callable, Awaitable
|
||||
from enum import Enum
|
||||
|
||||
|
||||
class TaskType(Enum):
|
||||
"""Types of agent sub-tasks."""
|
||||
TEST_ENDPOINT = "test_endpoint"
|
||||
VERIFY_FINDING = "verify_finding"
|
||||
SEARCH_CVE = "search_cve"
|
||||
GENERATE_POC = "generate_poc"
|
||||
CHAIN_EXPLORE = "chain_explore"
|
||||
DEEP_TEST = "deep_test"
|
||||
RECON_EXPAND = "recon_expand"
|
||||
BANNER_CHECK = "banner_check"
|
||||
MUTATE_TEST = "mutate_test"
|
||||
|
||||
|
||||
class TaskStatus(Enum):
|
||||
"""Task lifecycle states."""
|
||||
PENDING = "pending"
|
||||
RUNNING = "running"
|
||||
COMPLETED = "completed"
|
||||
FAILED = "failed"
|
||||
CANCELLED = "cancelled"
|
||||
|
||||
|
||||
# Priority levels (lower number = higher priority)
|
||||
PRIORITY_CRITICAL = 1 # RCE, auth bypass, chain exploitation
|
||||
PRIORITY_HIGH = 2 # Confirmed reflection, SQL error, SSRF indicators
|
||||
PRIORITY_MEDIUM = 3 # Standard vulnerability testing
|
||||
PRIORITY_LOW = 4 # Info disclosure, header checks
|
||||
PRIORITY_INFO = 5 # Enhancement, non-critical recon
|
||||
|
||||
|
||||
@dataclass
|
||||
class AgentTask:
|
||||
"""A discrete task that can be assigned to the agent or a specialist."""
|
||||
id: str = ""
|
||||
task_type: str = TaskType.TEST_ENDPOINT.value
|
||||
priority: int = PRIORITY_MEDIUM
|
||||
target: str = ""
|
||||
parameters: Dict = field(default_factory=dict)
|
||||
status: str = TaskStatus.PENDING.value
|
||||
result: Any = None
|
||||
error: str = ""
|
||||
created_at: float = 0.0
|
||||
started_at: float = 0.0
|
||||
completed_at: float = 0.0
|
||||
source: str = "" # What created this task (e.g., "chain_engine", "reasoning")
|
||||
|
||||
def __post_init__(self):
|
||||
if not self.id:
|
||||
self.id = uuid.uuid4().hex[:12]
|
||||
if not self.created_at:
|
||||
self.created_at = time.time()
|
||||
|
||||
def __lt__(self, other):
|
||||
"""For PriorityQueue comparison — lower priority number = higher priority."""
|
||||
if not isinstance(other, AgentTask):
|
||||
return NotImplemented
|
||||
return self.priority < other.priority
|
||||
|
||||
@property
|
||||
def age_seconds(self) -> float:
|
||||
return time.time() - self.created_at
|
||||
|
||||
@property
|
||||
def duration_seconds(self) -> Optional[float]:
|
||||
if self.started_at and self.completed_at:
|
||||
return self.completed_at - self.started_at
|
||||
return None
|
||||
|
||||
|
||||
class AgentTaskManager:
|
||||
"""Manages agent sub-tasks with priority queue and concurrent execution.
|
||||
|
||||
Tasks are submitted with priorities and processed concurrently.
|
||||
Supports task cancellation, status tracking, and result collection.
|
||||
"""
|
||||
|
||||
MAX_QUEUE_SIZE = 500
|
||||
MAX_COMPLETED = 200
|
||||
|
||||
def __init__(self, max_concurrent: int = 5):
|
||||
self.max_concurrent = max_concurrent
|
||||
self._queue: asyncio.PriorityQueue = asyncio.PriorityQueue(
|
||||
maxsize=self.MAX_QUEUE_SIZE
|
||||
)
|
||||
self._running: Dict[str, AgentTask] = {}
|
||||
self._completed: List[AgentTask] = []
|
||||
self._failed: List[AgentTask] = []
|
||||
self._cancelled = False
|
||||
self._semaphore = asyncio.Semaphore(max_concurrent)
|
||||
self._total_submitted = 0
|
||||
self._total_completed = 0
|
||||
self._total_failed = 0
|
||||
|
||||
# ── Task Submission ──
|
||||
|
||||
async def submit(self, task: AgentTask) -> str:
|
||||
"""Submit task to priority queue. Returns task ID."""
|
||||
if self._cancelled:
|
||||
return ""
|
||||
|
||||
if self._queue.full():
|
||||
# Evict lowest-priority pending task
|
||||
# Note: PriorityQueue doesn't support removal, so we skip
|
||||
return ""
|
||||
|
||||
try:
|
||||
self._queue.put_nowait(task)
|
||||
self._total_submitted += 1
|
||||
return task.id
|
||||
except asyncio.QueueFull:
|
||||
return ""
|
||||
|
||||
def submit_sync(self, task: AgentTask) -> str:
|
||||
"""Synchronous submit (for use in non-async contexts)."""
|
||||
try:
|
||||
self._queue.put_nowait(task)
|
||||
self._total_submitted += 1
|
||||
return task.id
|
||||
except (asyncio.QueueFull, Exception):
|
||||
return ""
|
||||
|
||||
async def submit_batch(self, tasks: List[AgentTask]) -> List[str]:
|
||||
"""Submit multiple tasks at once."""
|
||||
ids = []
|
||||
for task in tasks:
|
||||
tid = await self.submit(task)
|
||||
if tid:
|
||||
ids.append(tid)
|
||||
return ids
|
||||
|
||||
# ── Task Execution ──
|
||||
|
||||
async def run_tasks(self, executor: Callable[[AgentTask], Awaitable[Any]],
|
||||
cancel_check: Optional[Callable[[], bool]] = None,
|
||||
progress_callback: Optional[Callable[[Dict], Awaitable]] = None
|
||||
) -> List[AgentTask]:
|
||||
"""Process queue with concurrent execution.
|
||||
|
||||
Args:
|
||||
executor: Async function that executes a single task
|
||||
cancel_check: Optional function that returns True to stop
|
||||
progress_callback: Optional async callback for progress updates
|
||||
|
||||
Returns:
|
||||
List of completed tasks
|
||||
"""
|
||||
workers = []
|
||||
completed_in_run = []
|
||||
|
||||
async def worker():
|
||||
while not self._cancelled:
|
||||
if cancel_check and cancel_check():
|
||||
break
|
||||
|
||||
try:
|
||||
task = await asyncio.wait_for(
|
||||
self._queue.get(), timeout=2.0
|
||||
)
|
||||
except asyncio.TimeoutError:
|
||||
# Check if queue is permanently empty
|
||||
if self._queue.empty():
|
||||
break
|
||||
continue
|
||||
|
||||
if self._cancelled or (cancel_check and cancel_check()):
|
||||
break
|
||||
|
||||
async with self._semaphore:
|
||||
task.status = TaskStatus.RUNNING.value
|
||||
task.started_at = time.time()
|
||||
self._running[task.id] = task
|
||||
|
||||
try:
|
||||
result = await executor(task)
|
||||
task.result = result
|
||||
task.status = TaskStatus.COMPLETED.value
|
||||
task.completed_at = time.time()
|
||||
self._total_completed += 1
|
||||
completed_in_run.append(task)
|
||||
|
||||
# Bounded completed list
|
||||
self._completed.append(task)
|
||||
if len(self._completed) > self.MAX_COMPLETED:
|
||||
self._completed = self._completed[-self.MAX_COMPLETED:]
|
||||
|
||||
except Exception as e:
|
||||
task.error = str(e)
|
||||
task.status = TaskStatus.FAILED.value
|
||||
task.completed_at = time.time()
|
||||
self._total_failed += 1
|
||||
self._failed.append(task)
|
||||
|
||||
finally:
|
||||
self._running.pop(task.id, None)
|
||||
|
||||
if progress_callback:
|
||||
try:
|
||||
await progress_callback(self.get_status())
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# Spawn workers
|
||||
for _ in range(min(self.max_concurrent, max(1, self._queue.qsize()))):
|
||||
workers.append(asyncio.create_task(worker()))
|
||||
|
||||
if workers:
|
||||
await asyncio.gather(*workers, return_exceptions=True)
|
||||
|
||||
return completed_in_run
|
||||
|
||||
async def drain(self, executor: Callable[[AgentTask], Awaitable[Any]],
|
||||
cancel_check: Optional[Callable[[], bool]] = None,
|
||||
timeout: float = 300.0) -> List[AgentTask]:
|
||||
"""Run all queued tasks with a timeout."""
|
||||
try:
|
||||
return await asyncio.wait_for(
|
||||
self.run_tasks(executor, cancel_check),
|
||||
timeout=timeout,
|
||||
)
|
||||
except asyncio.TimeoutError:
|
||||
self.cancel()
|
||||
return list(self._completed[-50:])
|
||||
|
||||
# ── Task Control ──
|
||||
|
||||
def cancel(self):
|
||||
"""Cancel all pending tasks."""
|
||||
self._cancelled = True
|
||||
# Drain queue
|
||||
while not self._queue.empty():
|
||||
try:
|
||||
task = self._queue.get_nowait()
|
||||
task.status = TaskStatus.CANCELLED.value
|
||||
except asyncio.QueueEmpty:
|
||||
break
|
||||
|
||||
def reset(self):
|
||||
"""Reset for new use."""
|
||||
self._cancelled = False
|
||||
self._queue = asyncio.PriorityQueue(maxsize=self.MAX_QUEUE_SIZE)
|
||||
self._running.clear()
|
||||
|
||||
# ── Status & Queries ──
|
||||
|
||||
def get_status(self) -> Dict:
|
||||
"""Return task manager status for logging/dashboard."""
|
||||
return {
|
||||
"queued": self._queue.qsize(),
|
||||
"running": len(self._running),
|
||||
"completed": self._total_completed,
|
||||
"failed": self._total_failed,
|
||||
"total_submitted": self._total_submitted,
|
||||
"cancelled": self._cancelled,
|
||||
}
|
||||
|
||||
@property
|
||||
def is_empty(self) -> bool:
|
||||
return self._queue.empty() and not self._running
|
||||
|
||||
@property
|
||||
def pending_count(self) -> int:
|
||||
return self._queue.qsize()
|
||||
|
||||
@property
|
||||
def running_count(self) -> int:
|
||||
return len(self._running)
|
||||
|
||||
def get_completed_results(self, task_type: Optional[str] = None) -> List[Any]:
|
||||
"""Get results from completed tasks, optionally filtered by type."""
|
||||
tasks = self._completed
|
||||
if task_type:
|
||||
tasks = [t for t in tasks if t.task_type == task_type]
|
||||
return [t.result for t in tasks if t.result is not None]
|
||||
|
||||
def get_failed_tasks(self) -> List[AgentTask]:
|
||||
"""Get failed tasks for retry or debugging."""
|
||||
return list(self._failed[-50:])
|
||||
|
||||
|
||||
# ── Task Factory Helpers ──
|
||||
|
||||
def create_test_task(url: str, vuln_type: str, params: Dict = None,
|
||||
priority: int = PRIORITY_MEDIUM,
|
||||
source: str = "") -> AgentTask:
|
||||
"""Create a test_endpoint task."""
|
||||
return AgentTask(
|
||||
task_type=TaskType.TEST_ENDPOINT.value,
|
||||
priority=priority,
|
||||
target=url,
|
||||
parameters={
|
||||
"vuln_type": vuln_type,
|
||||
"params": params or {},
|
||||
},
|
||||
source=source,
|
||||
)
|
||||
|
||||
|
||||
def create_cve_task(software: str, version: str,
|
||||
priority: int = PRIORITY_HIGH) -> AgentTask:
|
||||
"""Create a search_cve task."""
|
||||
return AgentTask(
|
||||
task_type=TaskType.SEARCH_CVE.value,
|
||||
priority=priority,
|
||||
target=f"{software}/{version}",
|
||||
parameters={
|
||||
"software": software,
|
||||
"version": version,
|
||||
},
|
||||
source="cve_hunter",
|
||||
)
|
||||
|
||||
|
||||
def create_chain_task(finding_type: str, finding_url: str,
|
||||
chain_target: Dict,
|
||||
priority: int = PRIORITY_HIGH) -> AgentTask:
|
||||
"""Create a chain_explore task."""
|
||||
return AgentTask(
|
||||
task_type=TaskType.CHAIN_EXPLORE.value,
|
||||
priority=priority,
|
||||
target=finding_url,
|
||||
parameters={
|
||||
"source_vuln": finding_type,
|
||||
"chain_target": chain_target,
|
||||
},
|
||||
source="chain_engine",
|
||||
)
|
||||
|
||||
|
||||
def create_deep_test_task(url: str, vuln_type: str, param: str,
|
||||
mutation_context: Dict = None,
|
||||
priority: int = PRIORITY_HIGH) -> AgentTask:
|
||||
"""Create a deep_test task (re-test with mutations)."""
|
||||
return AgentTask(
|
||||
task_type=TaskType.DEEP_TEST.value,
|
||||
priority=priority,
|
||||
target=url,
|
||||
parameters={
|
||||
"vuln_type": vuln_type,
|
||||
"param": param,
|
||||
"mutation_context": mutation_context or {},
|
||||
},
|
||||
source="payload_mutator",
|
||||
)
|
||||
|
||||
|
||||
def create_poc_task(finding_id: str, finding_type: str,
|
||||
priority: int = PRIORITY_LOW) -> AgentTask:
|
||||
"""Create a generate_poc task."""
|
||||
return AgentTask(
|
||||
task_type=TaskType.GENERATE_POC.value,
|
||||
priority=priority,
|
||||
target=finding_id,
|
||||
parameters={"vuln_type": finding_type},
|
||||
source="exploit_generator",
|
||||
)
|
||||
Reference in New Issue
Block a user