memory_firewall.py

# 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),
        }