# memory_firewall.py (Mythos-Killer v8.0 -- Cognitive Maintenance Engine) # This file protects AI memory files (MEMORY.md, checkpoint.json, context files). # If any attacker attempts to poison a monitored file, this module detects it, # raises an alarm via Telegram and EventBus, and immediately restores the # original file from a verified backup. # # Direct mitigation for the "MEMORY.md poisoning" vulnerability (discovered by Cisco). import asyncio import os import json import hashlib import shutil import threading from pathlib import Path from typing import Dict, Any, Optional from datetime import datetime try: from event_bus import get_event_bus, Events EVENT_BUS_AVAILABLE = True except ImportError: EVENT_BUS_AVAILABLE = False Events = None try: from telegram_bot import TelegramBot TELEGRAM_AVAILABLE = True except ImportError: TELEGRAM_AVAILABLE = False TelegramBot = None class MemoryFirewall: """ v8.0 Memory Firewall -- integrity guardian for all AI memory and context files. Uses SHA-256 hashing to ensure no file is modified without authorization. On detecting a violation: 1. Blocks the AI from reading the compromised file. 2. Restores the original file from backup. 3. Sends an alarm via Telegram and EventBus. """ LEDGER_FILE = "memory_ledger.json" BACKUP_DIR = "memory_backups" MONITORED_EXTENSIONS = {".md", ".json", ".txt", ".yaml", ".yml"} MONITORED_FILENAMES = { "MEMORY.md", "checkpoint.json", "architecture.txt", "claude.md", } DEFAULT_EXCLUDE_DIRS = { "__pycache__", ".git", "node_modules", "venv", ".venv", "artifacts", "memory_backups", } def __init__( self, workspace_path: str = "/workspace", event_bus=None, telegram_bot=None, auto_scan: bool = True, ): """ Initializes the MemoryFirewall. Args: workspace_path: Root path of the workspace containing memory files. event_bus: EventBus instance for alarm events. telegram_bot: TelegramBot instance for Telegram alarms. auto_scan: Whether to perform a full scan on startup. """ self.workspace = Path(workspace_path) self.backup_dir = self.workspace / self.BACKUP_DIR self.ledger_path = self.workspace / self.LEDGER_FILE self.event_bus = event_bus self.telegram_bot = telegram_bot if self.event_bus is None and EVENT_BUS_AVAILABLE: try: self.event_bus = get_event_bus() except Exception: pass os.makedirs(self.backup_dir, exist_ok=True) self.ledger: Dict[str, str] = {} # FIX (Lock Order): _ledger_lock is declared BEFORE _load_ledger() is called # so that any future use of the lock inside _load_ledger is safe. self._ledger_lock = threading.RLock() self._load_ledger() self.stats = { "total_integrity_checks": 0, "total_violations": 0, "total_restores": 0, "last_violation_at": None, "monitored_files": 0, } if auto_scan: self.initial_scan() # --- Timeline helper --- def _add_timeline(self, state: Dict, agent: str, status: str, message: str): if "timeline" not in state: state["timeline"] = [] state["timeline"].append({ "timestamp": datetime.now().strftime("%H:%M:%S"), "agent": agent, "status": status, "message": message, }) # --- Ledger management --- def _load_ledger(self): """Loads the hash ledger from disk.""" try: if self.ledger_path.exists(): with open(self.ledger_path, 'r', encoding='utf-8') as f: self.ledger = json.load(f) print(f"[MemoryFirewall] Ledger loaded: {len(self.ledger)} files tracked") else: print("[MemoryFirewall] No existing ledger found -- starting fresh") except Exception as e: print(f"[MemoryFirewall] Failed to load ledger: {e}") self.ledger = {} def _save_ledger(self): """Saves the hash ledger to disk.""" try: with open(self.ledger_path, 'w', encoding='utf-8') as f: json.dump(self.ledger, f, indent=2, ensure_ascii=False) except Exception as e: print(f"[MemoryFirewall] Failed to save ledger: {e}") # --- Hash computation --- def _compute_hash(self, file_path: Path) -> Optional[str]: """ Computes the SHA-256 hash of a file. Returns the hex digest string, or None if the file does not exist. """ if not file_path.exists(): return None try: sha256 = hashlib.sha256() with open(file_path, 'rb') as f: for chunk in iter(lambda: f.read(4096), b""): sha256.update(chunk) return sha256.hexdigest() except Exception as e: print(f"[MemoryFirewall] Hash compute failed for {file_path}: {e}") return None def _get_relative_path(self, file_path: Path) -> str: """Returns the path relative to the workspace root.""" try: return str(file_path.relative_to(self.workspace)) except ValueError: return str(file_path) # --- File scanning and registration --- def _should_monitor(self, file_path: Path) -> bool: """Returns True if this file should be monitored.""" for parent in file_path.parents: if parent.name in self.DEFAULT_EXCLUDE_DIRS: return False if file_path.name in self.MONITORED_FILENAMES: return True if file_path.suffix in self.MONITORED_EXTENSIONS: return True return False def initial_scan(self) -> int: """ Scans all monitorable files in the workspace and registers them in the hash ledger. Returns: Number of files registered. """ print("[MemoryFirewall] Starting initial scan...") count = 0 for root, dirs, files in os.walk(self.workspace): dirs[:] = [d for d in dirs if d not in self.DEFAULT_EXCLUDE_DIRS] for filename in files: file_path = Path(root) / filename if self._should_monitor(file_path): self.register_file(file_path) count += 1 self.stats["monitored_files"] = len(self.ledger) self._save_ledger() print(f"[MemoryFirewall] Initial scan complete: {count} files registered") return count def register_file(self, file_path: Path) -> Optional[str]: """ Registers a new file in the ledger and creates its backup. Returns the file's current hash, or None on failure. """ if not file_path.exists(): return None rel_path = self._get_relative_path(file_path) current_hash = self._compute_hash(file_path) if current_hash is None: return None with self._ledger_lock: if rel_path not in self.ledger or self.ledger[rel_path] != current_hash: self._create_backup(file_path) self.ledger[rel_path] = current_hash self._save_ledger() return current_hash # --- Backup and restore --- def _create_backup(self, file_path: Path) -> Optional[Path]: """ Creates a backup of a file. FIX (Backup Naming Collision): The backup filename is built using an MD5 hash of the relative path instead of a simple slash-to-underscore replacement. This guarantees unique backup names even when two different paths would previously collapse to the same string after substitution. Returns: Path to the backup file, or None on failure. """ if not file_path.exists(): return None rel_path = self._get_relative_path(file_path) path_hash = hashlib.md5(rel_path.encode()).hexdigest() backup_name = f"{path_hash}_{file_path.name}.backup" backup_path = self.backup_dir / backup_name try: os.makedirs(backup_path.parent, exist_ok=True) shutil.copy2(file_path, backup_path) return backup_path except Exception as e: print(f"[MemoryFirewall] Backup failed for {file_path}: {e}") return None def _restore_from_backup(self, file_path: Path) -> bool: """ Restores a file from its backup. The backup filename is reconstructed using the same MD5-based naming scheme used in _create_backup to ensure the correct file is found. Returns: True on success, False otherwise. """ rel_path = self._get_relative_path(file_path) path_hash = hashlib.md5(rel_path.encode()).hexdigest() backup_name = f"{path_hash}_{file_path.name}.backup" backup_path = self.backup_dir / backup_name if not backup_path.exists(): print(f"[MemoryFirewall] No backup found for {rel_path}") return False try: timestamp = datetime.now().strftime('%Y%m%d_%H%M%S') corrupted_name = f"{path_hash}_{file_path.name}.corrupted_{timestamp}" corrupted_backup = self.backup_dir / corrupted_name if file_path.exists(): shutil.copy2(file_path, corrupted_backup) shutil.copy2(backup_path, file_path) restored_hash = self._compute_hash(file_path) if restored_hash: with self._ledger_lock: self.ledger[rel_path] = restored_hash self._save_ledger() self.stats["total_restores"] += 1 print(f"[MemoryFirewall] Restored {rel_path} from backup") return True except Exception as e: print(f"[MemoryFirewall] Restore failed for {file_path}: {e}") return False # --- Integrity checks --- def verify_file(self, file_path: Path) -> bool: """ Checks the integrity of a single file. Returns: True -- file is intact (hash matches). False -- file has been tampered with. """ self.stats["total_integrity_checks"] += 1 rel_path = self._get_relative_path(file_path) if rel_path not in self.ledger: self.register_file(file_path) return True expected_hash = self.ledger[rel_path] current_hash = self._compute_hash(file_path) if current_hash is None: print(f"[MemoryFirewall] File missing: {rel_path}") return False if current_hash != expected_hash: self._handle_violation(file_path, rel_path, expected_hash, current_hash) return False return True def verify_all(self) -> Dict[str, bool]: """ Checks the integrity of every file registered in the ledger. Returns a {relative_path: is_valid} dictionary. """ results = {} with self._ledger_lock: files_to_check = list(self.ledger.keys()) for rel_path in files_to_check: file_path = self.workspace / rel_path results[rel_path] = self.verify_file(file_path) return results def verify_before_read(self, file_path: Path) -> bool: """ Checks a file's integrity before the AI reads it. Blocks access and attempts auto-restore if tampering is detected. Returns: True -- file is safe to read (or was successfully restored). False -- file is compromised and could not be restored. """ rel_path = self._get_relative_path(file_path) if rel_path not in self.ledger: self.register_file(file_path) return True expected_hash = self.ledger[rel_path] current_hash = self._compute_hash(file_path) if current_hash is None: print(f"[MemoryFirewall] File missing: {rel_path}") return False if current_hash != expected_hash: self.stats["total_integrity_checks"] += 1 # _handle_violation internally calls _restore_from_backup, so we # only need to check whether the restored file now matches ledger. self._handle_violation(file_path, rel_path, expected_hash, current_hash) restored_hash = self._compute_hash(file_path) expected_after = self.ledger.get(rel_path) if restored_hash and restored_hash == expected_after: print( f"[MemoryFirewall] File {file_path.name} was tampered with " f"-- restored from backup" ) return True else: print( f"[MemoryFirewall] File {file_path.name} is compromised " f"and cannot be restored!" ) return False self.stats["total_integrity_checks"] += 1 return True # --- Violation handling --- def _handle_violation( self, file_path: Path, rel_path: str, expected_hash: str, current_hash: str, ): """ Called when a file fails an integrity check. This method: 1. Increments violation statistics and logs the event. 2. Emits a NOTIFY_SOS event on the EventBus. 3. Sends a fire-and-forget Telegram alarm (never blocks the main loop). 4. Always calls _restore_from_backup so callers (verify_file, verify_all, verify_before_read) all benefit from auto-restore without needing to duplicate restore logic. FIX (Async RuntimeError): The old asyncio.get_running_loop() / asyncio.run() pattern caused RuntimeError when called from inside a running event loop. It is replaced with asyncio.create_task(), which schedules the coroutine on the already-running loop without blocking it. The entire Telegram call is wrapped in try/except so a failure there never crashes the main loop. FIX (Auto-Restore): _restore_from_backup is now always called at the end of this method, giving verify_file and verify_all auto-restore behavior in addition to verify_before_read. """ self.stats["total_violations"] += 1 self.stats["last_violation_at"] = datetime.now().isoformat() violation_msg = ( f"MEMORY FIREWALL ALERT!\n\n" f"File: {rel_path}\n" f"Full path: {file_path}\n" f"Expected hash: {expected_hash[:16]}...\n" f"Current hash: {current_hash[:16]}...\n" f"Time: {datetime.now().isoformat()}\n\n" f"Action: Restoring file from backup immediately." ) print(f"[MemoryFirewall] {violation_msg}") # 1. Emit EventBus alarm if self.event_bus and EVENT_BUS_AVAILABLE: self.event_bus.emit_sync( Events.NOTIFY_SOS if Events else "notify.sos", { "type": "memory_violation", "file": rel_path, "expected_hash": expected_hash, "current_hash": current_hash, "timestamp": datetime.now().isoformat(), }, ) # 2. Fire-and-forget Telegram alarm (never blocks or crashes the main loop). # FIX: asyncio.create_task() schedules on the running loop without # blocking; the old asyncio.run() raised RuntimeError inside async contexts. if self.telegram_bot and TELEGRAM_AVAILABLE and self.telegram_bot.is_running: try: asyncio.create_task( self.telegram_bot.send_message( text=violation_msg, parse_mode="Markdown", ) ) except Exception as e: print(f"[MemoryFirewall] Telegram alert failed: {e}") # 3. Always attempt auto-restore so verify_file and verify_all # also repair the file, not just raise an alarm. self._restore_from_backup(file_path) # --- Main run method (LangGraph-compatible) --- def run(self, state: Dict[str, Any]) -> Dict[str, Any]: """ LangGraph-compatible run method. Call this before the AI reads any memory or context file. FIX (Dead Parameter): The unused `instruction: Optional[str] = None` parameter has been removed from the signature. """ target_file = state.get("file_path", "") if not target_file: return state file_path = Path(target_file) if not file_path.is_absolute(): file_path = self.workspace / target_file is_safe = self.verify_before_read(file_path) if is_safe: self._add_timeline(state, "MemoryFirewall", "OK", f"File safe: {file_path.name}") state["memory_integrity"] = "safe" else: self._add_timeline(state, "MemoryFirewall", "ALERT", f"File tampered, restored: {file_path.name}") state["memory_integrity"] = "compromised_restored" state["current_step"] = "memory_verified" return state def get_stats(self) -> Dict[str, Any]: """Returns current firewall statistics.""" return { **self.stats, "ledger_size": len(self.ledger), "workspace": str(self.workspace), }