mirror of
https://github.com/CyberSecurityUP/NeuroSploit.git
synced 2026-02-12 14:02:45 +00:00
220 lines
8.2 KiB
Python
220 lines
8.2 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
Scan Scheduler - Recurring task orchestration for NeuroSploit.
|
|
|
|
Supports cron expressions and interval-based scheduling for:
|
|
- Reconnaissance scans
|
|
- Vulnerability validation
|
|
- Re-analysis of previous findings
|
|
|
|
Uses APScheduler with SQLite persistence so jobs survive restarts.
|
|
"""
|
|
|
|
import json
|
|
import logging
|
|
from datetime import datetime
|
|
from typing import Dict, List, Optional
|
|
from pathlib import Path
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
try:
|
|
from apscheduler.schedulers.asyncio import AsyncIOScheduler
|
|
from apscheduler.triggers.cron import CronTrigger
|
|
from apscheduler.triggers.interval import IntervalTrigger
|
|
from apscheduler.jobstores.sqlalchemy import SQLAlchemyJobStore
|
|
HAS_APSCHEDULER = True
|
|
except ImportError:
|
|
HAS_APSCHEDULER = False
|
|
logger.warning("APScheduler not installed. Scheduler disabled. Install with: pip install apscheduler>=3.10.0")
|
|
|
|
|
|
class ScanScheduler:
|
|
"""Manages recurring scan jobs via APScheduler."""
|
|
|
|
def __init__(self, config: Dict, database_url: str = "sqlite:///./data/neurosploit_scheduler.db"):
|
|
self.config = config
|
|
self.scheduler_config = config.get('scheduler', {})
|
|
self.enabled = self.scheduler_config.get('enabled', False)
|
|
self.jobs_meta: Dict[str, Dict] = {} # job_id -> metadata
|
|
self._scan_callback = None
|
|
|
|
if not HAS_APSCHEDULER:
|
|
self.enabled = False
|
|
self.scheduler = None
|
|
return
|
|
|
|
jobstores = {
|
|
'default': SQLAlchemyJobStore(url=database_url)
|
|
}
|
|
self.scheduler = AsyncIOScheduler(jobstores=jobstores)
|
|
|
|
# Load pre-configured jobs from config
|
|
for job_config in self.scheduler_config.get('jobs', []):
|
|
try:
|
|
self.add_job(
|
|
job_id=job_config['id'],
|
|
target=job_config['target'],
|
|
scan_type=job_config.get('scan_type', 'quick'),
|
|
cron_expression=job_config.get('cron'),
|
|
interval_minutes=job_config.get('interval_minutes'),
|
|
agent_role=job_config.get('agent_role'),
|
|
llm_profile=job_config.get('llm_profile')
|
|
)
|
|
except Exception as e:
|
|
logger.error(f"Failed to load scheduled job '{job_config.get('id', '?')}': {e}")
|
|
|
|
def set_scan_callback(self, callback):
|
|
"""Set the callback function that executes scans.
|
|
|
|
The callback signature should be:
|
|
async def callback(target: str, scan_type: str,
|
|
agent_role: Optional[str], llm_profile: Optional[str]) -> Dict
|
|
"""
|
|
self._scan_callback = callback
|
|
|
|
def add_job(self, job_id: str, target: str, scan_type: str = "quick",
|
|
cron_expression: Optional[str] = None,
|
|
interval_minutes: Optional[int] = None,
|
|
agent_role: Optional[str] = None,
|
|
llm_profile: Optional[str] = None) -> Dict:
|
|
"""Schedule a recurring scan job.
|
|
|
|
Args:
|
|
job_id: Unique identifier for the job
|
|
target: Target URL or IP
|
|
scan_type: 'quick', 'full', 'recon', or 'analysis'
|
|
cron_expression: Cron schedule (e.g., '0 */6 * * *' for every 6 hours)
|
|
interval_minutes: Alternative to cron - run every N minutes
|
|
agent_role: Optional agent role for AI analysis
|
|
llm_profile: Optional LLM profile override
|
|
"""
|
|
if not self.scheduler:
|
|
return {"error": "Scheduler not available (APScheduler not installed)"}
|
|
|
|
if cron_expression:
|
|
trigger = CronTrigger.from_crontab(cron_expression)
|
|
schedule_desc = f"cron: {cron_expression}"
|
|
elif interval_minutes:
|
|
trigger = IntervalTrigger(minutes=interval_minutes)
|
|
schedule_desc = f"every {interval_minutes} minutes"
|
|
else:
|
|
return {"error": "Provide either cron_expression or interval_minutes"}
|
|
|
|
self.scheduler.add_job(
|
|
self._execute_scheduled_scan,
|
|
trigger=trigger,
|
|
id=job_id,
|
|
args=[target, scan_type, agent_role, llm_profile],
|
|
replace_existing=True,
|
|
name=f"scan_{target}_{scan_type}"
|
|
)
|
|
|
|
meta = {
|
|
"id": job_id,
|
|
"target": target,
|
|
"scan_type": scan_type,
|
|
"schedule": schedule_desc,
|
|
"agent_role": agent_role,
|
|
"llm_profile": llm_profile,
|
|
"created_at": datetime.now().isoformat(),
|
|
"last_run": None,
|
|
"run_count": 0,
|
|
"status": "active"
|
|
}
|
|
self.jobs_meta[job_id] = meta
|
|
|
|
logger.info(f"Scheduled job '{job_id}': {target} ({scan_type}) - {schedule_desc}")
|
|
return meta
|
|
|
|
def remove_job(self, job_id: str) -> bool:
|
|
"""Remove a scheduled job."""
|
|
if not self.scheduler:
|
|
return False
|
|
try:
|
|
self.scheduler.remove_job(job_id)
|
|
self.jobs_meta.pop(job_id, None)
|
|
logger.info(f"Removed scheduled job: {job_id}")
|
|
return True
|
|
except Exception as e:
|
|
logger.error(f"Failed to remove job '{job_id}': {e}")
|
|
return False
|
|
|
|
def pause_job(self, job_id: str) -> bool:
|
|
"""Pause a scheduled job."""
|
|
if not self.scheduler:
|
|
return False
|
|
try:
|
|
self.scheduler.pause_job(job_id)
|
|
if job_id in self.jobs_meta:
|
|
self.jobs_meta[job_id]["status"] = "paused"
|
|
return True
|
|
except Exception as e:
|
|
logger.error(f"Failed to pause job '{job_id}': {e}")
|
|
return False
|
|
|
|
def resume_job(self, job_id: str) -> bool:
|
|
"""Resume a paused job."""
|
|
if not self.scheduler:
|
|
return False
|
|
try:
|
|
self.scheduler.resume_job(job_id)
|
|
if job_id in self.jobs_meta:
|
|
self.jobs_meta[job_id]["status"] = "active"
|
|
return True
|
|
except Exception as e:
|
|
logger.error(f"Failed to resume job '{job_id}': {e}")
|
|
return False
|
|
|
|
def list_jobs(self) -> List[Dict]:
|
|
"""List all scheduled jobs with metadata."""
|
|
jobs = []
|
|
if self.scheduler:
|
|
for job in self.scheduler.get_jobs():
|
|
meta = self.jobs_meta.get(job.id, {})
|
|
jobs.append({
|
|
"id": job.id,
|
|
"name": job.name,
|
|
"next_run": str(job.next_run_time) if job.next_run_time else None,
|
|
"target": meta.get("target", "unknown"),
|
|
"scan_type": meta.get("scan_type", "unknown"),
|
|
"schedule": meta.get("schedule", "unknown"),
|
|
"status": meta.get("status", "active"),
|
|
"last_run": meta.get("last_run"),
|
|
"run_count": meta.get("run_count", 0)
|
|
})
|
|
return jobs
|
|
|
|
async def _execute_scheduled_scan(self, target: str, scan_type: str,
|
|
agent_role: Optional[str],
|
|
llm_profile: Optional[str]):
|
|
"""Execute a scheduled scan. Called by APScheduler."""
|
|
job_id = f"scan_{target}_{scan_type}"
|
|
logger.info(f"Executing scheduled scan: {target} ({scan_type})")
|
|
|
|
if job_id in self.jobs_meta:
|
|
self.jobs_meta[job_id]["last_run"] = datetime.now().isoformat()
|
|
self.jobs_meta[job_id]["run_count"] += 1
|
|
|
|
if self._scan_callback:
|
|
try:
|
|
result = await self._scan_callback(target, scan_type, agent_role, llm_profile)
|
|
logger.info(f"Scheduled scan completed: {target} ({scan_type})")
|
|
return result
|
|
except Exception as e:
|
|
logger.error(f"Scheduled scan failed for {target}: {e}")
|
|
else:
|
|
logger.warning("No scan callback registered. Scheduled scan skipped.")
|
|
|
|
def start(self):
|
|
"""Start the scheduler."""
|
|
if self.scheduler and self.enabled:
|
|
self.scheduler.start()
|
|
logger.info(f"Scheduler started with {len(self.list_jobs())} jobs")
|
|
|
|
def stop(self):
|
|
"""Stop the scheduler gracefully."""
|
|
if self.scheduler and self.scheduler.running:
|
|
self.scheduler.shutdown(wait=False)
|
|
logger.info("Scheduler stopped")
|