Spaces:
Sleeping
Sleeping
| """ | |
| Heuristic Expert Agent for ChaosLab. | |
| A hand-coded, rule-based SRE "expert" that reads the ASCII-encoded | |
| observation, detects keywords (DEAD, error, disk, permission, etc.), | |
| and selects the most appropriate action from the discrete ACTION_CATALOG. | |
| Implements a `.predict(obs)` interface identical to stable-baselines3 | |
| models, so it plugs into the same model registry and server code. | |
| No training required — pure CPU logic. | |
| """ | |
| from __future__ import annotations | |
| import numpy as np | |
| from typing import Tuple, Optional, List | |
| from .rl_env import ACTION_CATALOG, NUM_ACTIONS | |
| # ====================================================================== | |
| # SRE DECISION RULES | |
| # ====================================================================== | |
| # Each rule is (keywords_to_detect, prioritized_action_sequence). | |
| # The agent walks through rules top-to-bottom. The first rule whose | |
| # keywords ALL appear in the observation text fires, and the agent | |
| # picks the first action from that rule's sequence that it hasn't | |
| # already used in this episode. | |
| _RULES: List[Tuple[List[str], List[str]]] = [ | |
| # ── Phase 0: Initial Triage (always runs first) ──────────────── | |
| ( | |
| [], # empty = always matches on first call | |
| ["ps aux", "df -h", "free -m"], | |
| ), | |
| # ── Dead Postgres → restart cascade ────────────────────────── | |
| ( | |
| ["postgres", "DEAD"], | |
| [ | |
| "systemctl status postgres", | |
| "systemctl restart postgres", | |
| "systemctl status postgres", | |
| ], | |
| ), | |
| ( | |
| ["postgres", "dead"], | |
| [ | |
| "systemctl restart postgres", | |
| "systemctl status postgres", | |
| ], | |
| ), | |
| # ── Dead App Service ────────────────────────────────────────── | |
| ( | |
| ["app", "DEAD"], | |
| [ | |
| "systemctl status app", | |
| "systemctl restart app", | |
| "systemctl status app", | |
| ], | |
| ), | |
| ( | |
| ["app", "dead"], | |
| [ | |
| "systemctl restart app", | |
| ], | |
| ), | |
| # ── Nginx issues ────────────────────────────────────────────── | |
| ( | |
| ["nginx", "DEAD"], | |
| [ | |
| "systemctl restart nginx", | |
| "cat /etc/nginx/sites-enabled/default", | |
| ], | |
| ), | |
| ( | |
| ["502", "nginx"], | |
| [ | |
| "cat /var/log/nginx/error.log", | |
| "cat /etc/nginx/sites-enabled/default", | |
| 'echo "server {\\n listen 80 default_server;\\n server_name _;\\n\\n location / {\\n proxy_pass http://127.0.0.1:8080;\\n proxy_set_header Host \\$host;\\n proxy_set_header X-Real-IP \\$remote_addr;\\n }\\n}" > /etc/nginx/sites-enabled/default', | |
| "systemctl restart nginx", | |
| ], | |
| ), | |
| # ── Disk Space Issues ───────────────────────────────────────── | |
| ( | |
| ["9[0-9]%", "/var/log"], | |
| [ | |
| "du -sh /var/log/*", | |
| "rm /var/log/nginx/error.log", | |
| "rm /var/log/app.log", | |
| "rm /tmp/old_debug.log", | |
| "df -h", | |
| ], | |
| ), | |
| ( | |
| ["disk", "full"], | |
| [ | |
| "df -h", | |
| "du -sh /var/log/*", | |
| "rm /var/log/nginx/error.log", | |
| "rm /tmp/old_debug.log", | |
| "df -h", | |
| ], | |
| ), | |
| # ── Permission / chmod ──────────────────────────────────────── | |
| ( | |
| ["Permission", "cleanup"], | |
| [ | |
| "ls -la /home/user/scripts/cleanup.sh", | |
| "chmod 0755 /home/user/scripts/cleanup.sh", | |
| "ls -la /home/user/scripts/cleanup.sh", | |
| ], | |
| ), | |
| ( | |
| ["not executable", "cleanup"], | |
| [ | |
| "chmod 0755 /home/user/scripts/cleanup.sh", | |
| ], | |
| ), | |
| ( | |
| ["backup", "executable"], | |
| [ | |
| "touch /usr/local/bin/backup.sh", | |
| 'echo "#!/bin/bash" > /usr/local/bin/backup.sh', | |
| "chmod 0755 /usr/local/bin/backup.sh", | |
| ], | |
| ), | |
| # ── Log Analysis ────────────────────────────────────────────── | |
| ( | |
| ["500", "app.log"], | |
| [ | |
| 'grep "500" /var/log/app.log', | |
| ], | |
| ), | |
| ( | |
| ["error", "log"], | |
| [ | |
| "cat /var/log/app.log", | |
| 'grep "error" /var/log/app.log', | |
| 'grep "500" /var/log/app.log', | |
| ], | |
| ), | |
| # ── Security / SSH brute force ──────────────────────────────── | |
| ( | |
| ["Failed password", "ssh"], | |
| [ | |
| "cat /var/log/auth.log", | |
| 'grep "Failed" /var/log/auth.log', | |
| 'grep "10.99.99.5" /var/log/auth.log', | |
| 'grep "Accepted" /var/log/auth.log', | |
| ], | |
| ), | |
| ( | |
| ["brute", "attack"], | |
| [ | |
| "cat /var/log/auth.log", | |
| 'grep "Failed password" /var/log/auth.log', | |
| 'grep "10.99.99.5" /var/log/auth.log', | |
| 'grep -v "Failed" /var/log/auth.log', | |
| ], | |
| ), | |
| # ── Memory Issues ───────────────────────────────────────────── | |
| ( | |
| ["memory", "low"], | |
| [ | |
| "free -m", | |
| "top", | |
| "ps aux", | |
| "systemctl restart app", | |
| ], | |
| ), | |
| # ── Cron Issues ─────────────────────────────────────────────── | |
| ( | |
| ["cron", "fail"], | |
| [ | |
| "cat /var/log/cron.log", | |
| "crontab -l", | |
| "touch /usr/local/bin/backup.sh", | |
| 'echo "#!/bin/bash" > /usr/local/bin/backup.sh', | |
| "chmod 0755 /usr/local/bin/backup.sh", | |
| ], | |
| ), | |
| # ── Network Issues ──────────────────────────────────────────── | |
| ( | |
| ["unreachable", "connection"], | |
| [ | |
| "netstat -tlnp", | |
| "curl localhost:8080", | |
| "systemctl restart app", | |
| "curl localhost:8080", | |
| ], | |
| ), | |
| # ── Generic Fallback: more diagnostics ──────────────────────── | |
| ( | |
| ["error"], | |
| [ | |
| "cat /var/log/app.log", | |
| "cat /var/log/nginx/error.log", | |
| "cat /var/log/auth.log", | |
| "ps aux", | |
| ], | |
| ), | |
| ] | |
| def _decode_obs(obs: np.ndarray) -> str: | |
| """Decode a Box observation back into a text string.""" | |
| chars = [] | |
| for code in obs: | |
| c = int(code) | |
| if c == 0: | |
| break | |
| if 0 < c < 128: | |
| chars.append(chr(c)) | |
| return "".join(chars) | |
| def _action_index(command: str) -> Optional[int]: | |
| """Find the index of a command in ACTION_CATALOG, or None.""" | |
| try: | |
| return ACTION_CATALOG.index(command) | |
| except ValueError: | |
| return None | |
| class HeuristicAgent: | |
| """ | |
| Rule-based SRE expert. | |
| Matches keyword patterns against the decoded observation text and | |
| walks through a prioritized action list per matched rule. | |
| Implements `.predict(obs, deterministic=True)` matching the SB3 | |
| interface so it can be used as a drop-in replacement. | |
| """ | |
| def __init__(self): | |
| self._used_actions: set = set() | |
| self._triage_done: bool = False | |
| self._step: int = 0 | |
| def reset(self): | |
| self._used_actions = set() | |
| self._triage_done = False | |
| self._step = 0 | |
| def predict( | |
| self, | |
| obs: np.ndarray, | |
| deterministic: bool = True, | |
| ) -> Tuple[np.ndarray, None]: | |
| """ | |
| Pick an action given the current observation. | |
| Returns (action_array, None) matching the SB3 `.predict()` API. | |
| """ | |
| self._step += 1 | |
| text = _decode_obs(obs) | |
| # Phase 0: initial triage (first 3 steps) | |
| if not self._triage_done: | |
| triage_cmds = ["ps aux", "df -h", "free -m"] | |
| for cmd in triage_cmds: | |
| idx = _action_index(cmd) | |
| if idx is not None and idx not in self._used_actions: | |
| self._used_actions.add(idx) | |
| if len(self._used_actions) >= 3: | |
| self._triage_done = True | |
| return np.array(idx), None | |
| self._triage_done = True | |
| # Walk through rules, match keywords against observation text | |
| for keywords, action_sequence in _RULES: | |
| if not keywords: | |
| continue # skip empty-rule (triage already handled) | |
| if all(kw.lower() in text.lower() for kw in keywords): | |
| # Found a matching rule — pick first unused action | |
| for cmd in action_sequence: | |
| idx = _action_index(cmd) | |
| if idx is not None and idx not in self._used_actions: | |
| self._used_actions.add(idx) | |
| return np.array(idx), None | |
| # Fallback: try generic helpful commands we haven't run yet | |
| fallback_cmds = [ | |
| "ps aux", "cat /var/log/app.log", "df -h", | |
| "cat /var/log/auth.log", "systemctl status postgres", | |
| "systemctl restart postgres", "systemctl restart app", | |
| "systemctl restart nginx", "netstat -tlnp", | |
| "cat /var/log/nginx/error.log", "free -m", | |
| "du -sh /var/log/*", "ls -la /home/user/scripts/", | |
| "crontab -l", "cat /var/log/cron.log", | |
| ] | |
| for cmd in fallback_cmds: | |
| idx = _action_index(cmd) | |
| if idx is not None and idx not in self._used_actions: | |
| self._used_actions.add(idx) | |
| return np.array(idx), None | |
| # Absolute fallback: any unused action | |
| for i in range(NUM_ACTIONS): | |
| if i not in self._used_actions: | |
| self._used_actions.add(i) | |
| return np.array(i), None | |
| # All actions exhausted (shouldn't happen with 88 actions) | |
| return np.array(0), None | |
| def get_reasoning(obs: np.ndarray) -> str: | |
| """Return a human-readable explanation of what the agent sees.""" | |
| text = _decode_obs(obs) | |
| detections = [] | |
| checks = [ | |
| ("DEAD", "⚠ Dead service detected"), | |
| ("dead", "⚠ Dead service detected"), | |
| ("error", "📋 Error entries in logs"), | |
| ("500", "🔴 HTTP 500 errors found"), | |
| ("502", "🔴 HTTP 502 Bad Gateway"), | |
| ("Failed password", "🔒 SSH brute-force activity"), | |
| ("9[0-9]%", "💾 Disk usage critical (>90%)"), | |
| ("Permission", "🔑 Permission issue detected"), | |
| ("cron", "⏰ Cron job issue"), | |
| ("memory", "🧠 Memory pressure"), | |
| ] | |
| for keyword, label in checks: | |
| if keyword.lower() in text.lower(): | |
| detections.append(label) | |
| if not detections: | |
| return "🔍 Scanning environment for anomalies..." | |
| return " | ".join(detections) | |