mirror of
https://github.com/CyberSecurityUP/NeuroSploit.git
synced 2026-03-21 01:33:23 +00:00
e0935793c5
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
373 lines
12 KiB
Python
373 lines
12 KiB
Python
"""
|
|
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",
|
|
)
|