Add RLUtils class for managing RL/AI dashboard endpoints

- Implemented methods for fetching AI stats, training history, and recent experiences.
- Added functionality to set operation mode (MANUAL, AUTO, AI) with appropriate handling.
- Included helper methods for querying the database and sending JSON responses.
- Integrated model metadata extraction for visualization purposes.
This commit is contained in:
Fabien POLLY
2026-02-18 22:36:10 +01:00
parent b8a13cc698
commit eb20b168a6
684 changed files with 53278 additions and 27977 deletions

View File

@@ -3,11 +3,11 @@
"""
web_enum.py — Gobuster Web Enumeration -> DB writer for table `webenum`.
- Writes each finding into the `webenum` table
- ON CONFLICT(mac_address, ip, port, directory) DO UPDATE
- Respects orchestrator stop flag (shared_data.orchestrator_should_exit)
- No filesystem output: parse Gobuster stdout directly
- Filtrage dynamique des statuts HTTP via shared_data.web_status_codes
- Writes each finding into the `webenum` table in REAL-TIME (Streaming).
- Updates bjorn_progress with actual percentage (0-100%).
- Respects orchestrator stop flag (shared_data.orchestrator_should_exit) immediately.
- No filesystem output: parse Gobuster stdout/stderr directly.
- Filtrage dynamique des statuts HTTP via shared_data.web_status_codes.
"""
import re
@@ -15,6 +15,9 @@ import socket
import subprocess
import threading
import logging
import time
import os
import select
from typing import List, Dict, Tuple, Optional, Set
from shared import SharedData
@@ -27,8 +30,8 @@ b_class = "WebEnumeration"
b_module = "web_enum"
b_status = "WebEnumeration"
b_port = 80
b_service = '["http","https"]'
b_trigger = 'on_any:["on_web_service","on_new_port:80","on_new_port:443","on_new_port:8080","on_new_port:8443","on_new_port:9443","on_new_port:8000","on_new_port:8888","on_new_port:81","on_new_port:5000","on_new_port:5001","on_new_port:7080","on_new_port:9080"]'
b_service = '["http","https"]'
b_trigger = 'on_any:["on_web_service","on_new_port:80","on_new_port:443","on_new_port:8080","on_new_port:8443","on_new_port:9443","on_new_port:8000","on_new_port:8888","on_new_port:81","on_new_port:5000","on_new_port:5001","on_new_port:7080","on_new_port:9080"]'
b_parent = None
b_priority = 9
b_cooldown = 1800
@@ -36,8 +39,6 @@ b_rate_limit = '3/86400'
b_enabled = 1
# -------------------- Defaults & parsing --------------------
# Valeur de secours si l'UI n'a pas encore initialisé shared_data.web_status_codes
# (par défaut: 2xx utiles, 3xx, 401/403/405 et tous les 5xx; 429 non inclus)
DEFAULT_WEB_STATUS_CODES = [
200, 201, 202, 203, 204, 206,
301, 302, 303, 307, 308,
@@ -50,7 +51,6 @@ CTL_RE = re.compile(r"[\x00-\x1F\x7F]") # non-printables
# Gobuster "dir" line examples handled:
# /admin (Status: 301) [Size: 310] [--> http://10.0.0.5/admin/]
# /images (Status: 200) [Size: 12345]
GOBUSTER_LINE = re.compile(
r"""^(?P<path>\S+)\s*
\(Status:\s*(?P<status>\d{3})\)\s*
@@ -60,13 +60,14 @@ GOBUSTER_LINE = re.compile(
re.VERBOSE
)
# Regex pour capturer la progression de Gobuster sur stderr
# Ex: "Progress: 1024 / 4096 (25.00%)"
GOBUSTER_PROGRESS_RE = re.compile(r"Progress:\s+(?P<current>\d+)\s*/\s+(?P<total>\d+)")
def _normalize_status_policy(policy) -> Set[int]:
"""
Transforme une politique "UI" en set d'entiers HTTP.
policy peut contenir:
- int (ex: 200, 403)
- "xXX" (ex: "2xx", "5xx")
- "a-b" (ex: "500-504")
"""
codes: Set[int] = set()
if not policy:
@@ -99,30 +100,48 @@ def _normalize_status_policy(policy) -> Set[int]:
class WebEnumeration:
"""
Orchestrates Gobuster web dir enum and writes normalized results into DB.
In-memory only: no CSV, no temp files.
Streaming mode: Reads stdout/stderr in real-time for DB inserts and Progress UI.
"""
def __init__(self, shared_data: SharedData):
self.shared_data = shared_data
self.gobuster_path = "/usr/bin/gobuster" # verify with `which gobuster`
self.wordlist = self.shared_data.common_wordlist
self.lock = threading.Lock()
# Cache pour la taille de la wordlist (pour le calcul du %)
self.wordlist_size = 0
self._count_wordlist_lines()
# ---- Sanity checks
import os
self._available = True
if not os.path.exists(self.gobuster_path):
raise FileNotFoundError(f"Gobuster not found at {self.gobuster_path}")
logger.error(f"Gobuster not found at {self.gobuster_path}")
self._available = False
if not os.path.exists(self.wordlist):
raise FileNotFoundError(f"Wordlist not found: {self.wordlist}")
logger.error(f"Wordlist not found: {self.wordlist}")
self._available = False
# Politique venant de lUI : créer si absente
if not hasattr(self.shared_data, "web_status_codes") or not self.shared_data.web_status_codes:
self.shared_data.web_status_codes = DEFAULT_WEB_STATUS_CODES.copy()
logger.info(
f"WebEnumeration initialized (stdout mode, no files). "
f"Using status policy: {self.shared_data.web_status_codes}"
f"WebEnumeration initialized (Streaming Mode). "
f"Wordlist lines: {self.wordlist_size}. "
f"Policy: {self.shared_data.web_status_codes}"
)
def _count_wordlist_lines(self):
"""Compte les lignes de la wordlist une seule fois pour calculer le %."""
if self.wordlist and os.path.exists(self.wordlist):
try:
# Lecture rapide bufferisée
with open(self.wordlist, 'rb') as f:
self.wordlist_size = sum(1 for _ in f)
except Exception as e:
logger.error(f"Error counting wordlist lines: {e}")
self.wordlist_size = 0
# -------------------- Utilities --------------------
def _scheme_for_port(self, port: int) -> str:
https_ports = {443, 8443, 9443, 10443, 9444, 5000, 5001, 7080, 9080}
@@ -184,155 +203,195 @@ class WebEnumeration:
except Exception as e:
logger.error(f"DB insert error for {ip}:{port}{directory}: {e}")
# -------------------- Gobuster runner (stdout) --------------------
def _run_gobuster_stdout(self, url: str) -> Optional[str]:
base_cmd = [
self.gobuster_path, "dir",
"-u", url,
"-w", self.wordlist,
"-t", "10",
"--quiet",
"--no-color",
# Si supporté par ta version gobuster, tu peux réduire le bruit dès la source :
# "-b", "404,429",
]
def run(cmd):
return subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
# Try with -z first
cmd = base_cmd + ["-z"]
logger.info(f"Running Gobuster on {url}...")
try:
res = run(cmd)
if res.returncode == 0:
logger.success(f"Gobuster OK on {url}")
return res.stdout or ""
# Fallback if -z is unknown
if "unknown flag" in (res.stderr or "").lower() or "invalid" in (res.stderr or "").lower():
logger.info("Gobuster doesn't support -z, retrying without it.")
res2 = run(base_cmd)
if res2.returncode == 0:
logger.success(f"Gobuster OK on {url} (no -z)")
return res2.stdout or ""
logger.info(f"Gobuster failed on {url}: {res2.stderr.strip()}")
return None
logger.info(f"Gobuster failed on {url}: {res.stderr.strip()}")
return None
except Exception as e:
logger.error(f"Gobuster exception on {url}: {e}")
return None
def _parse_gobuster_text(self, text: str) -> List[Dict]:
"""
Parse gobuster stdout lines into entries:
{ 'path': '/admin', 'status': 301, 'size': 310, 'redirect': 'http://...'|None }
"""
entries: List[Dict] = []
if not text:
return entries
for raw in text.splitlines():
# 1) strip ANSI/control BEFORE regex
line = ANSI_RE.sub("", raw)
line = CTL_RE.sub("", line)
line = line.strip()
if not line:
continue
m = GOBUSTER_LINE.match(line)
if not m:
logger.debug(f"Unparsed line: {line}")
continue
# 2) extract all fields NOW
path = m.group("path") or ""
status = int(m.group("status"))
size = int(m.group("size") or 0)
redir = m.group("redir")
# 3) normalize path
if not path.startswith("/"):
path = "/" + path
path = "/" + path.strip("/")
entries.append({
"path": path,
"status": status,
"size": size,
"redirect": redir.strip() if redir else None
})
logger.info(f"Parsed {len(entries)} entries from gobuster stdout")
return entries
# -------------------- Public API --------------------
# -------------------- Public API (Streaming Version) --------------------
def execute(self, ip: str, port: int, row: Dict, status_key: str) -> str:
"""
Run gobuster on (ip,port), parse stdout, upsert each finding into DB.
Run gobuster on (ip,port), STREAM stdout/stderr, upsert findings real-time.
Updates bjorn_progress with 0-100% completion.
Returns: 'success' | 'failed' | 'interrupted'
"""
if not self._available:
return 'failed'
try:
if self.shared_data.orchestrator_should_exit:
logger.info("Interrupted before start (orchestrator flag).")
return "interrupted"
scheme = self._scheme_for_port(port)
base_url = f"{scheme}://{ip}:{port}"
logger.info(f"Enumerating {base_url} ...")
self.shared_data.bjornorch_status = "WebEnumeration"
if self.shared_data.orchestrator_should_exit:
logger.info("Interrupted before gobuster run.")
return "interrupted"
stdout_text = self._run_gobuster_stdout(base_url)
if stdout_text is None:
return "failed"
if self.shared_data.orchestrator_should_exit:
logger.info("Interrupted after gobuster run (stdout captured).")
return "interrupted"
entries = self._parse_gobuster_text(stdout_text)
if not entries:
logger.warning(f"No entries for {base_url}.")
return "success" # scan ran fine but no findings
# ---- Filtrage dynamique basé sur shared_data.web_status_codes
allowed = self._allowed_status_set()
pre = len(entries)
entries = [e for e in entries if e["status"] in allowed]
post = len(entries)
if post < pre:
preview = sorted(list(allowed))[:10]
logger.info(
f"Filtered out {pre - post} entries not in policy "
f"{preview}{'...' if len(allowed) > 10 else ''}."
)
# Setup Initial UI
self.shared_data.comment_params = {"ip": ip, "port": str(port), "url": base_url}
self.shared_data.bjorn_orch_status = "WebEnumeration"
self.shared_data.bjorn_progress = "0%"
logger.info(f"Enumerating {base_url} (Stream Mode)...")
# Prepare Identity & Policy
mac_address, hostname = self._extract_identity(row)
if not hostname:
hostname = self._reverse_dns(ip)
allowed = self._allowed_status_set()
for e in entries:
self._db_add_result(
mac_address=mac_address,
ip=ip,
hostname=hostname,
port=port,
directory=e["path"],
status=e["status"],
size=e.get("size", 0),
response_time=0, # gobuster doesn't expose timing here
content_type=None, # unknown here; a later HEAD/GET probe can fill it
tool="gobuster"
# Command Construction
# NOTE: Removed "--quiet" and "-z" to ensure we get Progress info on stderr
# But we use --no-color to make parsing easier
cmd = [
self.gobuster_path, "dir",
"-u", base_url,
"-w", self.wordlist,
"-t", "10", # Safe for RPi Zero
"--no-color",
"--no-progress=false", # Force progress bar even if redirected
]
process = None
findings_count = 0
stop_requested = False
# For progress calc
total_lines = self.wordlist_size if self.wordlist_size > 0 else 1
last_progress_update = 0
try:
# Merge stdout and stderr so we can read everything in one loop
process = subprocess.Popen(
cmd,
stdout=subprocess.PIPE,
stderr=subprocess.STDOUT,
text=True,
bufsize=1,
universal_newlines=True
)
return "success"
# Use select() (on Linux) so we can react quickly to stop requests
# without blocking forever on readline().
while True:
if self.shared_data.orchestrator_should_exit:
stop_requested = True
break
if process.poll() is not None:
# Process exited; drain remaining buffered output if any
line = process.stdout.readline() if process.stdout else ""
if not line:
break
else:
line = ""
if process.stdout:
if os.name != "nt":
r, _, _ = select.select([process.stdout], [], [], 0.2)
if r:
line = process.stdout.readline()
else:
# Windows: select() doesn't work on pipes; best-effort read.
line = process.stdout.readline()
if not line:
continue
# 3. Clean Line
clean_line = ANSI_RE.sub("", line).strip()
clean_line = CTL_RE.sub("", clean_line).strip()
if not clean_line:
continue
# 4. Check for Progress
if "Progress:" in clean_line:
now = time.time()
# Update UI max every 0.5s to save CPU
if now - last_progress_update > 0.5:
m_prog = GOBUSTER_PROGRESS_RE.search(clean_line)
if m_prog:
curr = int(m_prog.group("current"))
# Calculate %
pct = (curr / total_lines) * 100
pct = min(pct, 100.0)
self.shared_data.bjorn_progress = f"{int(pct)}%"
last_progress_update = now
continue
# 5. Check for Findings (Standard Gobuster Line)
m_res = GOBUSTER_LINE.match(clean_line)
if m_res:
st = int(m_res.group("status"))
# Apply Filtering Logic BEFORE DB
if st in allowed:
path = m_res.group("path")
if not path.startswith("/"): path = "/" + path
size = int(m_res.group("size") or 0)
redir = m_res.group("redir")
# Insert into DB Immediately
self._db_add_result(
mac_address=mac_address,
ip=ip,
hostname=hostname,
port=port,
directory=path,
status=st,
size=size,
response_time=0,
content_type=None,
tool="gobuster"
)
findings_count += 1
# Live feedback in comments
self.shared_data.comment_params = {
"url": base_url,
"found": str(findings_count),
"last": path
}
continue
# (Optional) Log errors/unknown lines if needed
# if "error" in clean_line.lower(): logger.debug(f"Gobuster err: {clean_line}")
# End of loop
if stop_requested:
logger.info("Interrupted by orchestrator.")
return "interrupted"
self.shared_data.bjorn_progress = "100%"
return "success"
except Exception as e:
logger.error(f"Execute error on {base_url}: {e}")
if process:
try:
process.terminate()
except Exception:
pass
return "failed"
finally:
if process:
try:
if stop_requested and process.poll() is None:
process.terminate()
# Always reap the child to avoid zombies.
try:
process.wait(timeout=2)
except Exception:
try:
process.kill()
except Exception:
pass
try:
process.wait(timeout=2)
except Exception:
pass
finally:
try:
if process.stdout:
process.stdout.close()
except Exception:
pass
self.shared_data.bjorn_progress = ""
self.shared_data.comment_params = {}
except Exception as e:
logger.error(f"Execute error on {ip}:{port}: {e}")
logger.error(f"General execution error: {e}")
return "failed"
@@ -341,7 +400,7 @@ if __name__ == "__main__":
shared_data = SharedData()
try:
web_enum = WebEnumeration(shared_data)
logger.info("Starting web directory enumeration...")
logger.info("Starting web directory enumeration (CLI)...")
rows = shared_data.read_data()
for row in rows:
@@ -351,6 +410,7 @@ if __name__ == "__main__":
port = row.get("port") or 80
logger.info(f"Execute WebEnumeration on {ip}:{port} ...")
status = web_enum.execute(ip, int(port), row, "enum_web_directories")
if status == "success":
logger.success(f"Enumeration successful for {ip}:{port}.")
elif status == "interrupted":