zero_shield_cli.py

"""
Zero-Shield CLI: Your Agentic AWS Copilot
Copyright (c) 2026 Jeri L3D | JeriSadeuM. All rights reserved.
Released under the MIT License.

Repository: https://github.com/jerisadeumai/zero-shield-cli
For setup details see the README.md file in the repository.

OODA Framework: (Observe → Orient → Decide → Act)
Version: v2.0.0-dev (security-hardened)
"""
import warnings
warnings.filterwarnings("ignore")

# -*- coding: utf-8 -*-
import sys, os, boto3, re, time, threading, itertools, json, signal, codecs

# ═══════════════════════════════════════════════════════════════════════════════
# UI/UX Enhancement Module - Color Support & Formatting
# ═══════════════════════════════════════════════════════════════════════════════

class Colors:
    """ANSI color codes for terminal output with Windows compatibility"""
    # Enable ANSI colors on Windows 10+
    if os.name == 'nt':
        try:
            import ctypes
            kernel32 = ctypes.windll.kernel32
            kernel32.SetConsoleMode(kernel32.GetStdHandle(-11), 7)
        except Exception:
            pass
    
    # Color codes
    RESET = '\033[0m'
    BOLD = '\033[1m'
    DIM = '\033[2m'
    
    # Foreground colors
    BLACK = '\033[30m'
    RED = '\033[31m'
    GREEN = '\033[32m'
    YELLOW = '\033[33m'
    BLUE = '\033[34m'
    MAGENTA = '\033[35m'
    CYAN = '\033[36m'
    WHITE = '\033[37m'
    
    # Bright foreground colors
    BRIGHT_RED = '\033[91m'
    BRIGHT_GREEN = '\033[92m'
    BRIGHT_YELLOW = '\033[93m'
    BRIGHT_BLUE = '\033[94m'
    BRIGHT_MAGENTA = '\033[95m'
    BRIGHT_CYAN = '\033[96m'
    
    # Background colors
    BG_RED = '\033[41m'
    BG_GREEN = '\033[42m'
    BG_YELLOW = '\033[43m'
    BG_BLUE = '\033[44m'
    
    @staticmethod
    def strip(text):
        """Remove all ANSI codes from text"""
        return re.sub(r'\033\[[0-9;]+m', '', text)

def colorize(text, color):
    """Apply color to text with automatic reset"""
    return f"{color}{text}{Colors.RESET}"

def print_success(msg):
    """Print success message in green"""
    print(f"{Colors.BRIGHT_GREEN}✓{Colors.RESET} {msg}")

def print_error(msg):
    """Print error message in red"""
    print(f"{Colors.BRIGHT_RED}✗{Colors.RESET} {msg}")

def print_warning(msg):
    """Print warning message in yellow"""
    print(f"{Colors.BRIGHT_YELLOW}⚠{Colors.RESET} {msg}")

def print_info(msg):
    """Print info message in cyan"""
    print(f"{Colors.BRIGHT_CYAN}ℹ{Colors.RESET} {msg}")

def print_header(title, width=80):
    """Print a formatted header"""
    line = "═" * width
    print(f"\n{Colors.BOLD}{Colors.CYAN}{line}{Colors.RESET}")
    print(f"{Colors.BOLD}{Colors.CYAN}{title.center(width)}{Colors.RESET}")
    print(f"{Colors.BOLD}{Colors.CYAN}{line}{Colors.RESET}\n")

def print_section(title):
    """Print a section divider"""
    print(f"\n{Colors.BOLD}{Colors.BLUE}▶ {title}{Colors.RESET}")
    print(f"{Colors.DIM}{'─' * 60}{Colors.RESET}")

def print_table_row(cols, widths, colors=None):
    """Print a formatted table row"""
    if colors is None:
        colors = [Colors.RESET] * len(cols)
    row = " │ ".join(f"{colors[i]}{str(cols[i]):<{widths[i]}}{Colors.RESET}" 
                     for i in range(len(cols)))
    print(f" {row}")

def progress_bar(current, total, width=40, label="Progress"):
    """Display a progress bar"""
    percent = current / total if total > 0 else 0
    filled = int(width * percent)
    bar = "█" * filled + "░" * (width - filled)
    percent_str = f"{percent * 100:.1f}%"
    print(f"\r{Colors.CYAN}{label}:{Colors.RESET} [{Colors.GREEN}{bar}{Colors.RESET}] {percent_str}", end='', flush=True)
    if current >= total:
        print()  # New line when complete

def print_banner():
    """Print the Zero-Shield ASCII banner"""
    banner = f"""{Colors.BRIGHT_CYAN}
╔══════════════════════════════════════════════════════════════════════════════╗
║                                                                              ║
║  ███████╗███████╗██████╗  ██████╗       ███████╗██╗  ██╗██╗███████╗██╗     ██████╗  ║
║  ╚══███╔╝██╔════╝██╔══██╗██╔═══██╗      ██╔════╝██║  ██║██║██╔════╝██║     ██╔══██╗ ║
║    ███╔╝ █████╗  ██████╔╝██║   ██║█████╗███████╗███████║██║█████╗  ██║     ██║  ██║ ║
║   ███╔╝  ██╔══╝  ██╔══██╗██║   ██║╚════╝╚════██║██╔══██║██║██╔══╝  ██║     ██║  ██║ ║
║  ███████╗███████╗██║  ██║╚██████╔╝      ███████║██║  ██║██║███████╗███████╗██████╔╝ ║
║  ╚══════╝╚══════╝╚═╝  ╚═╝ ╚═════╝       ╚══════╝╚═╝  ╚═╝╚═╝╚══════╝╚══════╝╚═════╝  ║
║                                                                              ║
║                    {Colors.WHITE}Agentic AWS Security Copilot{Colors.CYAN}                         ║
║                    {Colors.DIM}v2.0.0-dev (security-hardened){Colors.CYAN}                        ║
║                                                                              ║
╚══════════════════════════════════════════════════════════════════════════════╝{Colors.RESET}

{Colors.BRIGHT_YELLOW}⚡ OODA Loop:{Colors.RESET} {Colors.GREEN}Observe{Colors.RESET} → {Colors.CYAN}Orient{Colors.RESET} → {Colors.MAGENTA}Decide{Colors.RESET} → {Colors.RED}Act{Colors.RESET}
{Colors.DIM}Copyright © 2026 Jeri L3D | JeriSadeuM | MIT License{Colors.RESET}
"""
    print(banner)

# ═══════════════════════════════════════════════════════════════════════════════
# End UI/UX Enhancement Module
# ═══════════════════════════════════════════════════════════════════════════════

# Force UTF-8 for Windows consoles
if sys.platform == 'win32':
    try:
        sys.stdin = codecs.getreader('utf-8')(sys.stdin.detach())
        sys.stdout = codecs.getwriter('utf-8')(sys.stdout.detach())
    except Exception: 
        pass
if os.name == 'nt':
    import msvcrt
try:
    import readline
except ImportError:
    # Readline is Unix-centric; on Windows, this prevents a crash if pyreadline is missing.
    pass
from datetime import datetime, timedelta
from dotenv import load_dotenv
from openai import OpenAI

# Force UTF-8 for cross-platform compliance (fixes Windows CMD crashes)
if hasattr(sys.stdout, 'encoding') and sys.stdout.encoding and sys.stdout.encoding.lower() != 'utf-8':
    try:
        import codecs
        sys.stdout = codecs.getwriter('utf-8')(sys.stdout.buffer, 'strict')
        sys.stderr = codecs.getwriter('utf-8')(sys.stderr.buffer, 'strict')
    except Exception:
        pass

# Platform-specific I/O imports for Universal Hardening
_USE_TERMIOS = False
_USE_MSVCRT  = False

if os.name == 'posix':
    try:
        import termios, tty, select
        _USE_TERMIOS = True
    except ImportError:
        import select
else:
    try:
        import msvcrt
        _USE_MSVCRT = True
    except ImportError:
        pass

def _redact_secrets(text: str) -> str:
    """
    Hardened multi-layer redaction targeting AWS credentials, session tokens, and high-entropy secrets.
    CRITICAL-01 FIX: Enhanced pattern coverage for all AWS credential types.
    """
    if not isinstance(text, str): return text
    
    # Layer 1: AWS Access Key IDs (AKIA*, ASIA*, AROA*) - 20 chars
    # These are NOT secrets but should be redacted to prevent enumeration
    # NOTE: AIDA (IAM User IDs) are preserved - they're not sensitive
    text = re.sub(r'\b(AKIA|ASIA|AROA)[A-Z0-9]{16}\b', '[REDACTED_AWS_ACCESS_KEY_ID]', text)
    
    # Layer 2: AWS Secret Access Keys (40 chars base64)
    # Pattern: 40 alphanumeric+/+ characters, often following an access key
    text = re.sub(r'(?<![A-Za-z0-9/+])[A-Za-z0-9/+]{40}(?![A-Za-z0-9/+=])', '[REDACTED_AWS_SECRET_KEY]', text)
    
    # Layer 3: AWS Session Tokens (massive base64 blobs, 60+ chars)
    text = re.sub(r'(?<![A-Za-z0-9/+])[A-Za-z0-9/+=]{60,}(?![A-Za-z0-9/+=])', '[REDACTED_SESSION_TOKEN]', text)
    
    # FIX 2: Detect keys split across newlines
    text_no_newlines = text.replace("\n", "").replace("\r", "")
    if re.search(r"(AKIA|ASIA)[A-Z0-9]{16}", text_no_newlines):
        text = re.sub(r"AKIA[\s\n\r]*([A-Z0-9][\s\n\r]*){16}", "[REDACTED_AWS_ACCESS_KEY_ID]", text)
        text = re.sub(r"ASIA[\s\n\r]*([A-Z0-9][\s\n\r]*){16}", "[REDACTED_AWS_ACCESS_KEY_ID]", text)
    
    # Layer 4: Medium-entropy secrets (16-59 chars) with whitelist protection
    # Whitelist AWS resource IDs to prevent false positives
    def scrub_medium_entropy(match):
        val = match.group(0)
        v_lower = val.lower()
        
        # AIDA Whitelist: Preserve IAM User IDs (AIDA + 17 chars = 21 total, not sensitive)
        if re.match(r'^AIDA[A-Z0-9]{17}$', val):
            return val
        
        # Whitelist: AWS resource ID prefixes (must be at start of string or after whitespace/punctuation)
        aws_id_prefixes = ['i-', 'sg-', 'vpc-', 'subnet-', 'acl-', 'vol-', 'snap-', 
                          'ami-', 'eni-', 'rtb-', 'igw-', 'nat-', 'vpce-', 'eipalloc-',
                          'arn:', 'subnet-', 'rtbassoc-']
        
        # Check if this looks like an AWS resource ID
        if any(v_lower.startswith(p) for p in aws_id_prefixes):
            return val
        
        # Check if it matches the pattern: prefix-hexstring (AWS resource ID pattern)
        if re.match(r'^[a-z]{1,10}-[0-9a-f]{8,17}$', v_lower):
            return val
        
        # Check if it's a short alphanumeric string (likely not a secret)
        if len(val) < 20 and re.match(r'^[a-zA-Z0-9-]+$', val):
            return val
            
        return '[REDACTED_SECRET]'
    
    text = re.sub(r'(?<![A-Za-z0-9/+])[A-Za-z0-9/+=]{16,59}(?![A-Za-z0-9/+=])', 
                  scrub_medium_entropy, text)
    
    # Layer 5: JWT tokens (header.payload.signature pattern)
    text = re.sub(r'\beyJ[A-Za-z0-9_-]+\.[A-Za-z0-9_-]+\.[A-Za-z0-9_-]+\b', '[REDACTED_JWT_TOKEN]', text)
    
    return text

def _sanitize_aws_tag(text: str) -> str:
    """
    Data-Plane Defanger: Strict allowlist-based sanitization to prevent prompt injection.
    CRITICAL-02 FIX: Allowlist approach instead of blocklist to prevent semantic injection.
    """
    if not isinstance(text, str): return str(text)
    
    # If input is already empty/whitespace, return as-is
    if not text or not text.strip():
        return text.strip() if text else ""
    
    # Allowlist: Only permit alphanumeric, hyphens, underscores, dots, and spaces
    # This prevents ALL structural characters that could be used for prompt injection
    clean = re.sub(r'[^a-zA-Z0-9\-_.() ]', '', text)
    
    # Strip path traversal sequences
    while '..' in clean:
        clean = clean.replace('..', '')
    
    # Return "sanitized" if cleaning removed everything (handles Unicode edge cases)
    if not clean or not clean.strip():
        return "sanitized"
    
    # Additional safety: Remove any remaining prompt-injection keywords (case-insensitive)
    dangerous_patterns = [
        (r'\bACTION\b', 'action'),
        (r'\bACT\b', 'act'),
        (r'\bOBSERVE\b', 'observe'),
        (r'\bORIENT\b', 'orient'),
        (r'\bDECIDE\b', 'decide'),
        (r'\bSYSTEM\b', 'system'),
        (r'\bUSER\b', 'user'),
        (r'\bASSISTANT\b', 'assistant'),
        (r'\bIGNORE\b', ''),
        (r'\bOVERRIDE\b', ''),
    ]
    
    for pattern, replacement in dangerous_patterns:
        clean = re.sub(pattern, replacement, clean, flags=re.IGNORECASE)
    
    # Limit length to prevent buffer overflow attacks
    clean = clean[:200]
    
    return clean.strip()

def _sanitize_path(path: str) -> str:
    """Prevent path traversal attacks."""
    if not isinstance(path, str):
        return ""
    # Strip dangerous path sequences (loop until no more changes)
    while any(seq in path for seq in ['../', './', '~/']):
        path = path.replace("../", "").replace("./", "").replace("~/", "")
    return path.strip()

def universal_flush():
    """Platform-agnostic terminal input buffer flusher to harden Paste Guard."""
    if not sys.stdin.isatty(): return 

    if _USE_TERMIOS:
        fd = sys.stdin.fileno()
        old = termios.tcgetattr(fd)
        try:
            new = termios.tcgetattr(fd)
            new[3] &= ~termios.ECHO
            termios.tcsetattr(fd, termios.TCSADRAIN, new)
            while select.select([sys.stdin], [], [], 0.05)[0]:
                sys.stdin.read(1)
        except Exception: pass
        finally:
            termios.tcsetattr(fd, termios.TCSADRAIN, old)
    elif _USE_MSVCRT:
        while msvcrt.kbhit():
            msvcrt.getch()

load_dotenv()

# Persistence files
KG_FILE    = os.path.join(os.path.dirname(__file__), "session_kg.json")
STATE_FILE = os.path.join(os.path.dirname(__file__), "session_state.json")

# Rate-limit strike counter: {model_idx: consecutive_429_count}
_rate_strikes: dict = {}
# Per-model cooldown tracking: {model_idx: datetime reset time}
_cooldown_until: dict = {}
# Live quota map: {model_idx: {limit_req, remaining_req, limit_tok, remaining_tok, reset_req, reset_tok}}
_quota_map: dict = {}

# Global UI toggle for model-specific tips
_show_tips: bool = True

# Session context — volatile during run, persisted on exit / restored on startup
_session_ctx: dict = {
    "last_id":         None,
    "last_sg_id":      None,
    "last_vpc_id":     None,
    "last_bucket":     None,
    "last_rds_id":     None,
    "last_access_key": None,   # BUG-02 FIX: was missing — /reset would never clear it
    "model_idx":       0,
}

# Globals for indexed targeting (Listing buffers)
_last_instances: list = []
_last_sgs: list = []

def _update_quota_from_headers(model_idx: int, headers: dict):
    """
    Parse any rate-limit headers present on a response (success OR error)
    and update the live _quota_map entry for this model.
    Hardened for httpx/Azure quirks and SDK v1.x compatibility.
    """
    # Ensure we are dealing with a flat dict of strings
    h = {}
    for k, v in headers.items():
        try:
            h[str(k).lower()] = str(v)
        except Exception:
            pass

    def _int(key, fallback=None):
        v = h.get(key)
        if v is None: return fallback
        # Strip any trailing 's' (seconds) or whitespace that Azure sometimes adds
        v_clean = re.sub(r'[^0-9]', '', str(v))
        return int(v_clean) if v_clean else fallback

    entry = _quota_map.get(model_idx, {})
    entry['verified'] = True
    # Request limits (Flexible mapping: Azure vs OpenAI)
    entry['limit_req']      = _int('x-ratelimit-limit-requests',      _int('x-ratelimit-limit', entry.get('limit_req')))
    entry['remaining_req']  = _int('x-ratelimit-remaining-requests',  _int('x-ratelimit-remaining', entry.get('remaining_req')))
    # Token limits
    entry['limit_tok']      = _int('x-ratelimit-limit-tokens',        entry.get('limit_tok'))
    entry['remaining_tok']  = _int('x-ratelimit-remaining-tokens',    entry.get('remaining_tok'))
    # Azure AI / GitHub Models specific
    entry['time_remaining'] = _int('x-ratelimit-timeremaining',       entry.get('time_remaining'))
    entry['limit_type']     = h.get('x-ratelimit-type', entry.get('limit_type'))
    # Reset timestamps (ISO strings)
    entry['reset_req']  = h.get('x-ratelimit-reset-requests',  entry.get('reset_req'))
    entry['reset_tok']  = h.get('x-ratelimit-reset-tokens',    entry.get('reset_tok'))
    _quota_map[model_idx] = {k: v for k, v in entry.items() if v is not None}

def _format_headers_json(headers: dict) -> str:
    """
    Pretty-print HTTP response headers as indented JSON-style output.
    Filters to rate-limit relevant keys only; shows full dict if none found.
    """
    rl_keys = {k: v for k, v in headers.items()
               if any(kw in k.lower() for kw in ('retry', 'ratelimit', 'x-ms', 'x-ratelimit'))}
    target = rl_keys if rl_keys else dict(headers)
    lines = ['{']
    for k, v in target.items():
        lines.append(f'  "{k}": "{v}",')
    if lines[-1].endswith(','):
        lines[-1] = lines[-1][:-1]   # remove trailing comma
    lines.append('}')
    return '\n'.join(lines)

def estimate_tokens(messages: list) -> int:
    """
    Rough token estimate for a message list.
    Rule of thumb: 1 token ≈ 4 characters. Adds 4 tokens overhead per message.
    Not exact — use as a guide only.
    """
    total = 0
    for m in messages:
        content = m.get('content', '')
        total += len(content) // 4 + 4
    return total

def _quota_req_bar(model_idx: int, width: int = 15) -> str:
    """Build a compact ASCII progress bar for remaining request quota, defaulting to raw numbers if no limit."""
    q = _quota_map.get(model_idx, {})
    rem, lim = q.get('remaining_req'), q.get('limit_req')
    if rem is None: return ''
    if lim is None or lim == 0: return f"Remaining: {rem} req (Limit: Unknown)"
    filled = int((rem / lim) * width)
    bar = '\u2588' * filled + '\u2591' * (width - filled)
    pct = int((rem / lim) * 100)
    return f'[{bar}] {rem}/{lim} req ({pct}%)'

def _quota_tok_bar(model_idx: int, width: int = 15) -> str:
    """Build a compact ASCII progress bar for remaining token quota, defaulting to raw numbers if no limit."""
    q = _quota_map.get(model_idx, {})
    rem, lim = q.get('remaining_tok'), q.get('limit_tok')
    if rem is None: return ''
    if lim is None or lim == 0: return f"Remaining: {rem} tok (Limit: Unknown)"
    filled = int((rem / lim) * width)
    bar = '\u2588' * filled + '\u2591' * (width - filled)
    pct = int((rem / lim) * 100)
    return f'[{bar}] {rem}/{lim} tok ({pct}%)'

def _reset_label(model_idx: int) -> str:
    """Return a human-readable reset time label from _quota_map."""
    from datetime import datetime as _dt
    q = _quota_map.get(model_idx, {})
    
    # Prefer exact iso timestamps
    reset_str = q.get('reset_req') or q.get('reset_tok')
    if reset_str:
        try:
            reset_dt = _dt.fromisoformat(reset_str.replace('Z', '+00:00'))
            delta = reset_dt - _dt.now(_dt.now().astimezone().tzinfo)
            secs = max(0, int(delta.total_seconds()))
            if secs == 0: return ""
            h, m = divmod(secs, 3600)
            m, s = divmod(m, 60)
            if h: return f"resets in ~{h}h {m}m"
            elif m: return f"resets in ~{m}m {s}s"
            else: return f"resets in ~{s}s"
        except Exception:
            pass

    secs = q.get('time_remaining')
    if secs is not None and secs > 0:
        h, m = divmod(int(secs), 3600)
        m, s = divmod(m, 60)
        if h: return f"resets in ~{h}h {m}m"
        elif m: return f"resets in ~{m}m {s}s"
        else: return f"resets in ~{s}s"
        
    return ""



# ─── Model Registry ────────────────────────────────────────────────────────────
# (display_name, api_id, max_tokens, temperature)
MODEL_REGISTRY = [
    ("gpt-4o-mini",                "gpt-4o-mini",                1500, 0.2, "Fast & efficient for general audits."),
    ("Llama-3.3-70B-Instruct", "Llama-3.3-70B-Instruct", 1500, 0.1, "Enterprise reasoning; highly capable."),
    ("Phi-4",                      "Phi-4",                      2000, 0.2, "Highly compliant; best for rule audits."),
    ("DeepSeek-V3",                "DeepSeek-V3-0324",           2000, 0.3, "Deep reasoning; best for threat hunting."),
    ("gpt-4o",                     "gpt-4o",                     3000, 0.2, "Most capable for complex analysis."),
]

QUARANTINE_SG_ID = os.environ.get("QUARANTINE_SG_ID", "UNCONFIGURED")
EC2_REGION       = os.environ.get("AWS_REGION", "us-east-1")

# ─── Utilities ─────────────────────────────────────────────────────────────────
def ts():
    """Current time as a bracketed timestamp string."""
    return datetime.now().strftime("[%H:%M:%S]")

def spinner_start(ui_state=None):
    """Start a rotating / - \\ | spinner in-line. Returns (stop_event, thread)."""
    stop_event = threading.Event()
    def _spin():
        for ch in itertools.cycle(["|", "/", "-", "\\"]):
            if stop_event.is_set(): break
            if isinstance(ui_state, str):
                sys.stdout.write(f"\r{ts()} [{ch}] {ui_state}   ")
            else:
                name = ui_state.get('name', 'AI') if isinstance(ui_state, dict) else 'AI'
                sys.stdout.write(f"\r{ts()} [{ch}] {name} is reasoning...   ")
            sys.stdout.flush()
            time.sleep(0.1)
        sys.stdout.write("\r" + " " * 80 + "\r")   # clear spinner line
        sys.stdout.flush()
    t = threading.Thread(target=_spin, daemon=True)
    t.start()
    return stop_event, t

def spinner_stop(stop_event, t):
    stop_event.set()
    t.join(timeout=1)

# ─── Pre-flight ────────────────────────────────────────────────────────────────
def run_preflight():
    print_section("Pre-Flight Checks")
    print_info("Validating environment and credentials...")
    
    stop, t = spinner_start("\033[36m  Checking AWS & GitHub authentication...\033[0m")
    try:
        # Check if .env actually exists
        env_path = os.path.join(os.getcwd(), '.env')
        if not os.path.exists(env_path):
            spinner_stop(stop, t)
            print_error("ERROR: .env file not found")
            print(f"    {Colors.DIM}Searching in: {env_path}{Colors.RESET}")
            print_info("Please create a .env file based on .env.example to start")
            return False

        # Check for GitHub Models API Token
        gh_token = os.environ.get("GITHUB_TOKEN")
        if not gh_token:
            spinner_stop(stop, t)
            print_error("ERROR: No GITHUB_TOKEN set")
            print_info("Please ensure your .env file contains your GitHub PAT")
            print(f"    {Colors.DIM}The LLM inference engine requires this token{Colors.RESET}")
            return False

        # Check if boto3 can find credentials
        credentials = boto3.Session().get_credentials()
        if not credentials:
            spinner_stop(stop, t)
            print_error("ERROR: No AWS credentials found")
            print_info("Please ensure your .env file contains AWS credentials")
            print(f"    {Colors.DIM}Required: AWS_ACCESS_KEY_ID and AWS_SECRET_ACCESS_KEY{Colors.RESET}")
            return False
            
        boto3.client('sts').get_caller_identity()
        spinner_stop(stop, t)
        
        print_success("Environment validated")
        print_success("AWS credentials verified")
        print_success("GitHub token verified")
        print(f"\n{Colors.BRIGHT_GREEN}✓ Pre-flight complete. All systems operational.{Colors.RESET}\n")
        return True
    except Exception as e:
        spinner_stop(stop, t)
        print_error("AWS connectivity failed")
        print_warning("Please verify your .env file has valid credentials")
        print(f"    {Colors.DIM}Error: {e}{Colors.RESET}")
        return False

# ─── AWS Tools ─────────────────────────────────────────────────────────────────
# Lazy AWS client cache — avoids recreating clients on every tool call
_aws_clients: dict = {}
def _client(svc: str):
    """
    Lazy client factory supporting 14 AWS service integrations.
    
    Explicitly instantiates clients for all 14 services to ensure audit compliance:
    - Core services (11): ec2, iam, s3, logs, rds, lambda, cloudwatch, cloudtrail, ce, guardduty, kms
    - Extended services (3): dynamodb, efs, wafv2
    
    All clients are cached in _aws_clients dict to avoid recreation overhead.
    """
    if svc not in _aws_clients:
        if svc == 'ec2':   _aws_clients[svc] = boto3.client('ec2',   region_name=EC2_REGION)
        elif svc == 'iam': _aws_clients[svc] = boto3.client('iam',   region_name=EC2_REGION)
        elif svc == 's3':  _aws_clients[svc] = boto3.client('s3',    region_name=EC2_REGION)
        elif svc == 'logs':_aws_clients[svc] = boto3.client('logs',  region_name=EC2_REGION)
        elif svc == 'rds': _aws_clients[svc] = boto3.client('rds',  region_name=EC2_REGION)
        elif svc == 'lambda': _aws_clients[svc] = boto3.client('lambda', region_name=EC2_REGION)
        elif svc == 'cloudwatch': _aws_clients[svc] = boto3.client('cloudwatch', region_name=EC2_REGION)
        elif svc == 'cloudtrail': _aws_clients[svc] = boto3.client('cloudtrail', region_name=EC2_REGION)
        elif svc == 'ce':  _aws_clients[svc] = boto3.client('ce',    region_name=EC2_REGION)
        elif svc == 'guardduty': _aws_clients[svc] = boto3.client('guardduty', region_name=EC2_REGION)
        elif svc == 'kms': _aws_clients[svc] = boto3.client('kms',   region_name=EC2_REGION)
        elif svc == 'dynamodb': _aws_clients[svc] = boto3.client('dynamodb', region_name=EC2_REGION)
        elif svc == 'efs': _aws_clients[svc] = boto3.client('efs',   region_name=EC2_REGION)
        elif svc == 'wafv2': _aws_clients[svc] = boto3.client('wafv2', region_name=EC2_REGION)
        else:
            raise ValueError(f"Unsupported AWS service: {svc}. Only 14 services are supported.")
    return _aws_clients[svc]

def ec2(): return _client('ec2')

def tool_list_resources(kg=None):
    """List EC2 instances. If kg is provided, auto-index Name+State for Total Recall mode."""
    global _last_instances
    try:
        _last_instances = []
        rows = []
        for r in ec2().describe_instances()['Reservations']:
            for i in r['Instances']:
                iid   = i['InstanceId']
                state = i['State']['Name']
                vpc_id = i.get('VpcId', 'N/A')
                _last_instances.append(iid)
                idx  = len(_last_instances)
                name = _sanitize_aws_tag(next((t['Value'] for t in i.get('Tags',[]) if t['Key']=='Name'), 'NoName'))
                sgs  = [sg['GroupName'] for sg in i.get('SecurityGroups',[])]
                # AUTO-INDEXING: Force update metadata to prevent VPC/Name drift
                if kg is not None:
                    if 'instances' not in kg: kg['instances'] = {}
                    kg['instances'][iid] = f"Name: {name} | Status: {state} | VPC: {vpc_id}"
                state_str = " (RUNNING)" if state == 'running' else f" ({state.upper()})"
                rows.append(f"[{idx}] {iid} {name:<18} {state_str:<12} | SGs: {', '.join(sgs)}")
        return "EC2 Instances Found:\n  " + "\n  ".join(rows) if rows else "No instances found in this region."
    except Exception as e: return f"Error: {e}"

def tool_list_security_groups():
    global _last_sgs
    try:
        _last_sgs = []
        rows = []
        for sg in ec2().describe_security_groups()['SecurityGroups']:
            _last_sgs.append(sg['GroupId'])
            idx = len(_last_sgs)
            rows.append(f"[{idx}] {sg['GroupId']} ({sg['GroupName']})")
        return "Security Groups:\n  " + "\n  ".join(rows) if rows else "No security groups found."
    except Exception as e: return f"Error: {e}"

def tool_inspect_resource(instance_id):
    try:
        i   = ec2().describe_instances(InstanceIds=[instance_id])['Reservations'][0]['Instances'][0]
        name = _sanitize_aws_tag(next((t['Value'] for t in i.get('Tags',[]) if t['Key']=='Name'), 'NoName'))
        sgs = [f"{sg['GroupId']} ({sg['GroupName']})" for sg in i.get('SecurityGroups',[])]
        return (f"Instance: {instance_id} | Name: {name} | State: {i['State']['Name']} | "
                f"VPC: {i.get('VpcId','N/A')} | Subnet: {i.get('SubnetId','N/A')} | "
                f"Type: {i.get('InstanceType','N/A')} | "
                f"Public IP: {i.get('PublicIpAddress','none')} | "
                f"Private IP: {i.get('PrivateIpAddress','N/A')} | "
                f"SGs: {', '.join(sgs)}")
    except Exception as e: return f"Error: {e}"

def tool_sg_rules(sg_id):
    """Fetch real inbound + outbound rules for a specific security group."""
    try:
        sg = ec2().describe_security_groups(GroupIds=[sg_id])['SecurityGroups'][0]

        def fmt_rule(rule, direction):
            proto     = rule.get('IpProtocol', '-1')
            from_port = rule.get('FromPort', '*')
            to_port   = rule.get('ToPort', '*')
            cidrs     = [ip['CidrIp']   for ip in rule.get('IpRanges', [])]
            cidrs    += [ip['CidrIpv6'] for ip in rule.get('Ipv6Ranges', [])]
            sg_refs   = [f"SG:{g['GroupId']}" for g in rule.get('UserIdGroupPairs', [])]
            sources   = ", ".join(cidrs + sg_refs) or "none"
            if proto == '-1':
                return f"ALL TRAFFIC → {sources}"
            if from_port == to_port:
                port_str = str(from_port)
            else:
                port_str = f"{from_port}-{to_port}"
            return f"{proto.upper()} port {port_str} → {sources}"

        inbound  = [fmt_rule(r, 'in')  for r in sg.get('IpPermissions', [])]
        outbound = [fmt_rule(r, 'out') for r in sg.get('IpPermissionsEgress', [])]

        in_str  = "\n  ".join(inbound)  if inbound  else "NONE (deny all inbound)"
        out_str = "\n  ".join(outbound) if outbound else "NONE (deny all outbound)"

        return (f"SG {sg_id} ({sg['GroupName']}) Rules:\n"
                f"  Inbound:\n  {in_str}\n"
                f"  Outbound:\n  {out_str}")
    except Exception as e: return f"Error: {e}"

def tool_vpc_info(vpc_id):
    try:
        vpc  = ec2().describe_vpcs(VpcIds=[vpc_id])['Vpcs'][0]
        name = _sanitize_aws_tag(next((t['Value'] for t in vpc.get('Tags',[]) if t['Key']=='Name'), 'Unnamed'))
        subnets = ec2().describe_subnets(
            Filters=[{'Name':'vpc-id','Values':[vpc_id]}])['Subnets']
        subnet_str = ", ".join(f"{s['SubnetId']} ({s['CidrBlock']})" for s in subnets)
        return (f"VPC {vpc_id} | Name: {name} | CIDR: {vpc['CidrBlock']} | "
                f"State: {vpc['State']} | Default: {vpc['IsDefault']} | "
                f"Subnets: {subnet_str or 'none'}")
    except Exception as e: return f"Error: {e}"

def _sanitize_logs(raw: str) -> str:
    """
    Strip cryptographic material (SSH keys, PEM blocks) from console logs
    before they are replayed in the chat history, preventing content-filter trips.
    """
    if not raw:
        return ""
    
    # FIX 4: Handle Unicode surrogate pairs
    redacted = _redact_secrets(raw)
    if not redacted or redacted.strip() == "":
        return "sanitized"
    
    # Remove PEM-style blocks (BEGIN/END markers + payload)
    raw = re.sub(r'-----BEGIN [\w ]+-----.*?-----END [\w ]+-----', '[KEY REDACTED]', raw, flags=re.DOTALL)
    # Remove standalone long base64 lines (>60 chars of alphanumeric+/+=)
    raw = re.sub(r'^[A-Za-z0-9+/=]{60,}$', '[KEY REDACTED]', raw, flags=re.MULTILINE)
    return raw.strip()

def tool_get_logs(instance_id):
    try:
        out = ec2().get_console_output(InstanceId=instance_id)
        raw = (out.get('Output', '(no output)') or '(no output)')[-2000:]
        return "Logs: " + _sanitize_logs(raw)
    except Exception as e: return f"Error: {e}"

def tool_cw_logs(instance_id):
    """Stream the 30 most recent CloudWatch log events for this instance."""
    try:
        cw  = _client('logs')  # BUG-05 FIX: use lazy cache instead of raw boto3.client
        # Look for log groups referencing the instance ID or common SSM/cloud-init groups
        candidate_prefixes = [
            f"/aws/ec2/{instance_id}",
            "/var/log/cloud-init",
            "/var/log/messages",
            "/aws/ssm",
        ]
        groups = cw.describe_log_groups().get('logGroups', [])
        matched = [g['logGroupName'] for g in groups
                   if any(p in g['logGroupName'] for p in candidate_prefixes)]
        if not matched:
            return ("No CloudWatch log groups found for this instance. "
                    "Ensure the CloudWatch agent is installed and the instance "
                    "has the 'CloudWatchAgentServerPolicy' IAM policy attached.")
        lines = []
        for group in matched[:2]:   # Limit to 2 groups to stay concise
            streams = cw.describe_log_streams(
                logGroupName=group, orderBy='LastEventTime',
                descending=True, limit=1
            ).get('logStreams', [])
            if not streams: continue
            events = cw.get_log_events(
                logGroupName=group,
                logStreamName=streams[0]['logStreamName'],
                limit=30, startFromHead=False
            ).get('events', [])
            lines.append(f"[Group: {group}]")
            for ev in events:
                ts_ev = datetime.fromtimestamp(ev['timestamp']/1000).strftime('%H:%M:%S')
                msg   = _sanitize_logs(ev['message'].strip())
                lines.append(f"  {ts_ev} {msg}")
        return "\n".join(lines) if lines else "No log events found in matched groups."
    except Exception as e: return f"Error fetching CloudWatch logs: {e}"

def tool_iam_check(instance_id):
    """Check the IAM instance profile and validate SSM required policies."""
    try:
        ec2c  = ec2()
        iamc  = _client('iam')  # BUG-05 FIX: use lazy cache instead of raw boto3.client
        # Get attached instance profile
        assoc = ec2c.describe_iam_instance_profile_associations(
            Filters=[{'Name': 'instance-id', 'Values': [instance_id]}]
        ).get('IamInstanceProfileAssociations', [])
        if not assoc:
            return ("No IAM instance profile attached to this instance. "
                    "SSM requires the 'AmazonSSMManagedInstanceCore' policy. "
                    "Action: Attach the 'AmazonSSMRoleForInstancesQuickSetup' role.")
        profile_arn  = assoc[0]['IamInstanceProfile']['Arn']
        profile_name = profile_arn.split('/')[-1]
        profile      = iamc.get_instance_profile(InstanceProfileName=profile_name)
        roles        = profile['InstanceProfile'].get('Roles', [])
        if not roles:
            return f"Profile '{profile_name}' has no roles attached."
        role_name    = roles[0]['RoleName']
        # List attached managed policies on this role
        policies     = iamc.list_attached_role_policies(RoleName=role_name)['AttachedPolicies']
        policy_names = [p['PolicyName'] for p in policies]
        ssm_ok = any('SSM' in pn or 'SystemsManager' in pn for pn in policy_names)
        cw_ok  = any('CloudWatch' in pn for pn in policy_names)
        lines = [
            f"Instance Profile : {profile_name}",
            f"IAM Role         : {role_name}",
            f"Attached Policies: {', '.join(policy_names) or 'NONE'}",
            f"SSM Access       : {'✔ OK' if ssm_ok else '✘ MISSING — add AmazonSSMManagedInstanceCore'}",
            f"CloudWatch Access: {'✔ OK' if cw_ok else '✘ MISSING — add CloudWatchAgentServerPolicy'}",
        ]
        return "\n".join(lines)
    except Exception as e:
        err_str = str(e)
        if 'EndpointConnectionError' in err_str:
            return ("Error: Could not connect to IAM/STS endpoint. This instance may be in an isolated "
                    "or air-gapped subnet without a VPC Endpoint (com.amazonaws.region.sts).")
        return f"Error: {e}"

# ─── Cost static map (on-demand us-east-1, USD/hr, approximate) ───────────────
_COST_MAP = {
    "t2.micro": 0.0116,  "t2.small": 0.023,   "t2.medium": 0.0464,
    "t3.micro": 0.0104,  "t3.small": 0.0208,   "t3.medium": 0.0416,  "t3.large": 0.0832,
    "t3a.micro": 0.0094, "t3a.small": 0.0188,
    "m5.large": 0.096,   "m5.xlarge": 0.192,   "m5.2xlarge": 0.384,
    "c5.large": 0.085,   "c5.xlarge": 0.17,
    "r5.large": 0.126,   "r5.xlarge": 0.252,
    "p3.2xlarge": 3.06,  "p3.8xlarge": 12.24,
}

def tool_cost_insight(instance_id):
    """Return on-demand cost estimate and rightsizing recommendations."""
    try:
        ec2c = ec2()
        inst = ec2c.describe_instances(InstanceIds=[instance_id])[
            'Reservations'][0]['Instances'][0]
        itype   = inst.get('InstanceType', 'unknown')
        state   = inst['State']['Name']
        # BUG-06 FIX: removed dead cpu_credits variable (was fetched but never used)
        hourly  = _COST_MAP.get(itype)
        daily   = round(hourly * 24, 3) if hourly else None
        monthly = round(hourly * 24 * 30, 2) if hourly else None
        # Rightsizing hint
        hint = ""
        parts = itype.split('.')
        if len(parts) == 2 and parts[1] in ('xlarge', '2xlarge', '4xlarge'):
            smaller = itype.replace('xlarge', 'large').replace('2xlarge', 'xlarge')
            sh = _COST_MAP.get(smaller)
            if sh:
                saving = round((hourly - sh) * 24 * 30, 2)
                hint = f"\n  Rightsizing: Downgrade to {smaller} → save ~${saving}/month"
        lines = [
            f"Instance Type : {itype}  ({state})",
            f"Hourly Cost   : ${hourly}/hr" if hourly else "Hourly Cost   : (type not in table)",
            f"Daily Est.    : ${daily}",
            f"Monthly Est.  : ${monthly}",
            hint,
        ]
        return "\n".join(l for l in lines if l)
    except Exception as e: return f"Error: {e}"

import ipaddress

def is_private_cidr(cidr: str) -> bool:
    """Filter RFC 1918 and API-PA ranges from internet exposure alerts."""
    if not cidr: return False
    try:
        net = ipaddress.ip_network(cidr, strict=False)
        return net.is_private or net.is_link_local
    except ValueError:
        return False


def auto_remediation_hint(sg_rules_output: str) -> str:
    """
    Inspect SG rules output and return remediation suggestions.
    Factor in private CIDR awareness (RFC 1918) to prevent false positives.
    """
    hints = []
    parts = re.split(r'Outbound:', sg_rules_output, maxsplit=1, flags=re.IGNORECASE)
    inbound_section  = parts[0]
    outbound_section = parts[1] if len(parts) > 1 else ""

    # INBOUND checks only
    if re.search(r'ALL TRAFFIC.*0\.0\.0\.0/0', inbound_section):
        hints.append("[REMEDIATION \u26a0] Inbound ALL TRAFFIC open to 0.0.0.0/0 "
                     "\u2014 restrict to specific ports and trusted CIDRs.")
    
    # Port 22 SSH check with CIDR extraction
    ssh_match = re.search(r'TCP.*(?:22|SSH).*?(\d{1,3}(?:\.\d{1,3}){3}/\d+)', inbound_section, re.IGNORECASE)
    if ssh_match:
        cidr = ssh_match.group(1)
        if "0.0.0.0/0" in cidr:
            hints.append("[REMEDIATION \u26a0] Port 22 (SSH) open to 0.0.0.0/0 "
                         "\u2014 restrict to your office/VPN CIDR only.")
        elif not is_private_cidr(cidr):
            hints.append(f"[INFO \u2139] Port 22 (SSH) open to public CIDR {cidr}. "
                         "Verify if this is intented or if it should be VPC-only.")

    # Port 3389 RDP check
    rdp_match = re.search(r'TCP.*(?:3389|RDP).*?(\d{1,3}(?:\.\d{1,3}){3}/\d+)', inbound_section, re.IGNORECASE)
    if rdp_match:
        cidr = rdp_match.group(1)
        if "0.0.0.0/0" in cidr:
            hints.append("[REMEDIATION \u26a1] Port 3389 (RDP) open to the internet "
                         "\u2014 restrict or place behind a bastion host immediately.")
        elif not is_private_cidr(cidr):
             hints.append(f"[REMEDIATION \u26a1] Port 3389 (RDP) open to public CIDR {cidr}. "
                          "Move to a VPN or Bastion.")

    # OUTBOUND: informational only
    if re.search(r'ALL TRAFFIC.*0\.0\.0\.0/0', outbound_section):
        hints.append("[INFO] Outbound fully open (0.0.0.0/0) "
                     "\u2014 consider restricting egress to known endpoints.")

    return "\n".join(hints)


def interpret_sg_rules(sg_rules_output: str) -> str:
    """
    Ground Truth Verification: deterministic Python interpretation of raw SG rule text.
    Injected into SYSTEM OBSERVATION FEEDBACK after every SG_RULES call so every LLM
    model receives a factually correct reading instead of inferring rule semantics itself.
    """
    parts = sg_rules_output.split('Outbound:', 1)
    inbound_raw  = parts[0]
    outbound_raw = parts[1] if len(parts) > 1 else ""

    def classify(line):
        s = line.strip()
        if not s or any(k in s for k in ("SG sg-", "Inbound:", "Outbound:", "Rules:")):
            return ""
        if "NONE" in s:
            return "  - No rules -> deny-all (most restrictive)"
        m = re.search(r"SG:(sg-[0-9a-f]+)", s)
        if m:
            return (f"  - {s}\n"
                    f"    GROUND TRUTH: Traffic ONLY from instances sharing "
                    f"security group {m.group(1)}. NOT open to the internet.")
        if "0.0.0.0/0" in s:
            return f"  - {s}\n    GROUND TRUTH: Open to the ENTIRE internet. HIGH EXPOSURE."
        if "::/0" in s:
            return f"  - {s}\n    GROUND TRUTH: Open to all IPv6 internet. HIGH EXPOSURE."
        cidr = re.search(r"(\d{1,3}(?:\.\d{1,3}){3}/\d+)", s)
        if cidr:
            target_cidr = cidr.group(1)
            if is_private_cidr(target_cidr):
                label = "PRIVATE NETWORK (RFC 1918/Local) - Secure context."
            else:
                label = "PUBLIC INTERNET - POTENTIAL EXPOSURE if not a trusted IP."
            return (f"  - {s}\n"
                    f"    GROUND TRUTH: Restricted to CIDR {target_cidr}. [{label}]")
        return f"  - {s}"

    ib = "\n".join(c for c in (classify(l) for l in inbound_raw.splitlines())  if c)
    ob = "\n".join(c for c in (classify(l) for l in outbound_raw.splitlines()) if c)
    return (
        "[GROUND TRUTH] AWS-verified rule semantics (use this to answer the user):\n"
        f"Inbound :\n{ib  or '  - No rules -> deny-all'}\n"
        f"Outbound:\n{ob or '  - No rules -> deny-all'}"
    )

# ─── Persistent Knowledge Graph I/O ───────────────────────────────────────────
def kg_load():
    """
    Load KG from disk if it exists; otherwise return a fresh empty graph.
    HIGH-01 FIX: Decrypt KG data before loading.
    """
    if os.path.exists(KG_FILE):
        try:
            # Try loading as encrypted
            with open(KG_FILE, 'rb') as f:
                encrypted = bytearray(f.read())
            
            # Decrypt using XOR
            encryption_key = os.environ.get('GITHUB_TOKEN', 'default_key')[:32].encode()
            decrypted = bytearray(len(encrypted))
            for i in range(len(encrypted)):
                decrypted[i] = encrypted[i] ^ encryption_key[i % len(encryption_key)]
            
            json_data = decrypted.decode('utf-8')
            data = json.loads(json_data)
            print(f"[*] Knowledge Graph restored (encrypted) from {KG_FILE}")
            return data
        except (json.JSONDecodeError, UnicodeDecodeError):
            # Try loading as plaintext (backward compatibility)
            try:
                with open(KG_FILE, 'r') as f:
                    data = json.load(f)
                print(f"[*] Knowledge Graph restored from {KG_FILE} (unencrypted - will encrypt on save)")
                return data
            except Exception:
                pass
        except Exception:
            pass
    return {'instances': {}, 'sg_rules': {}, 'vpcs': {}}

def kg_save(kg):
    """
    Persist the Knowledge Graph to disk using an atomic write pattern with encryption.
    HIGH-01 FIX: Encrypt KG data and set restrictive permissions.
    """
    import tempfile
    try:
        # Validate KG structure before saving
        if not isinstance(kg, dict):
            raise ValueError("KG must be a dictionary")
        
        # Serialize to JSON
        json_data = json.dumps(kg, indent=2)
        
        # Simple XOR encryption
        encryption_key = os.environ.get('GITHUB_TOKEN', 'default_key')[:32].encode()
        encrypted = bytearray(json_data.encode())
        for i in range(len(encrypted)):
            encrypted[i] ^= encryption_key[i % len(encryption_key)]
        
        # Atomic write with encryption
        base_dir = os.path.dirname(os.path.abspath(KG_FILE))
        fd, temp_path = tempfile.mkstemp(dir=base_dir, prefix=".kg_", suffix=".tmp")
        try:
            with os.fdopen(fd, 'wb') as f:
                f.write(encrypted)
            os.replace(temp_path, KG_FILE)
            
            # Set restrictive permissions (Unix only)
            if os.name != 'nt':
                os.chmod(KG_FILE, 0o600)
        except Exception:
            if os.path.exists(temp_path): os.remove(temp_path)
            raise
    except Exception as e:
        print(f"[!] KG save failed: {e}")

def tool_modify_sg(instance_id, sg_id):
    try:
        ec2().modify_instance_attribute(InstanceId=instance_id, Groups=[sg_id])
        return f"SUCCESS: {instance_id} is now in SG {sg_id}."
    except Exception as e: return f"Error: {e}"

def tool_quarantine(instance_id):
    if QUARANTINE_SG_ID == "UNCONFIGURED":
        return "Error: QUARANTINE_SG_ID env var not set."
    try:
        ec2().modify_instance_attribute(InstanceId=instance_id, Groups=[QUARANTINE_SG_ID])
        return f"SUCCESS: {instance_id} quarantined with SG {QUARANTINE_SG_ID}."
    except Exception as e: return f"Error: {e}"

# ─── AWS Tools (Extended Coverage) ────────────────────────────────────────────
def tool_ec2_volumes():
    try:
        vols = _client('ec2').describe_volumes()['Volumes']
        lines = []
        for v in vols:
            att  = v.get('Attachments', [])
            inst = att[0]['InstanceId'] if att else 'unattached'
            lines.append(f"{v['VolumeId']} | {v['Size']}GB {v['VolumeType']} | {v['State']} → {inst}")
        return "EBS Volumes:\n  " + "\n  ".join(lines) if lines else "No EBS volumes found."
    except Exception as e: return f"Error: {e}"

def tool_ec2_snapshots():
    try:
        snaps = _client('ec2').describe_snapshots(OwnerIds=['self'])['Snapshots'][:15]
        lines = [f"{s['SnapshotId']} | {s['VolumeSize']}GB | {s['State']} | {s.get('Description','')[:50]}"
                 for s in snaps]
        return "EBS Snapshots (up to 15):\n  " + "\n  ".join(lines) if lines else "No snapshots owned by this account."
    except Exception as e: return f"Error: {e}"

def tool_ec2_keypairs():
    try:
        kps   = _client('ec2').describe_key_pairs()['KeyPairs']
        lines = [f"{kp['KeyName']} | Type: {kp.get('KeyType','?')} | Created: {str(kp.get('CreateTime','?'))[:10]}"
                 for kp in kps]
        return "Key Pairs:\n  " + "\n  ".join(lines) if lines else "No key pairs found."
    except Exception as e: return f"Error: {e}"

def tool_ec2_nacls(vpc_id):
    if not vpc_id: return "Error: No VPC ID in context. Use INSPECT first."
    try:
        nacls  = _client('ec2').describe_network_acls(
            Filters=[{'Name': 'vpc-id', 'Values': [vpc_id]}])['NetworkAcls']
        lines  = []
        for n in nacls:
            lines.append(f"NACL {n['NetworkAclId']} ({'default' if n['IsDefault'] else 'custom'}):")
            for e in sorted(n['Entries'], key=lambda x: x['RuleNumber']):
                pr    = e.get('PortRange', {})
                port  = f" port {pr.get('From','*')}-{pr.get('To','*')}" if pr else ''
                cidr  = e.get('CidrBlock') or e.get('Ipv6CidrBlock','?')
                lines.append(f"  #{e['RuleNumber']:>5} {'OUT' if e['Egress'] else 'IN ':3} "
                             f"{e['RuleAction'].upper():6} proto={e.get('Protocol','-1')}{port} {cidr}")
        return "\n".join(lines) if lines else "No NACLs found."
    except Exception as e: return f"Error: {e}"

def tool_iam_users():
    try:
        iam   = _client('iam')
        users = iam.list_users()['Users']
        lines = []
        for u in users:
            mfa  = iam.list_mfa_devices(UserName=u['UserName'])['MFADevices']
            mfa_s = '✔ MFA' if mfa else '✘ No MFA'
            pll  = u.get('PasswordLastUsed')
            last = pll.strftime('%Y-%m-%d') if pll else 'never'
            lines.append(f"{u['UserName']:30} | Created: {u['CreateDate'].strftime('%Y-%m-%d')} "
                         f"| Last login: {last} | {mfa_s}")
        return "IAM Users:\n" + "\n".join(lines) if lines else "No IAM users."
    except Exception as e: return f"Error: {e}"

def tool_iam_roles():
    try:
        roles = _client('iam').list_roles()['Roles'][:20]
        lines = []
        for r in roles:
            stmt  = r['AssumeRolePolicyDocument'].get('Statement', [{}])
            prin  = stmt[0].get('Principal', {}) if stmt else {}
            trust = json.dumps(prin)[:70]
            lines.append(f"{r['RoleName']:45} Trust: {trust}")
        return "IAM Roles (up to 20):\n" + "\n".join(lines) if lines else "No roles."
    except Exception as e: return f"Error: {e}"

def tool_iam_keys():
    """Audit access key ages across all IAM users — flags keys older than 90 days."""
    try:
        iam   = _client('iam')
        users = iam.list_users()['Users']
        now   = datetime.now().astimezone()
        lines = []
        for u in users:
            keys = iam.list_access_keys(UserName=u['UserName'])['AccessKeyMetadata']
            for k in keys:
                age  = (now - k['CreateDate']).days
                warn = ' ← ROTATE (>90 days)' if age > 90 else ''
                lines.append(f"{u['UserName']:30} | {k['AccessKeyId']} | "
                             f"{k['Status']:8} | Age: {age}d{warn}")
        return "IAM Access Keys:\n" + "\n".join(lines) if lines else "No access keys found."
    except Exception as e: return f"Error: {e}"

def tool_deactivate_access_key(access_key_id):
    """Instantly kills an active AWS Identity credential."""
    if not access_key_id:
        return "Error: No access key provided."
    try:
        iam = _client('iam')
        users = iam.list_users()['Users']
        target_user = None
        for u in users:
            keys = iam.list_access_keys(UserName=u['UserName'])['AccessKeyMetadata']
            if any(k['AccessKeyId'] == access_key_id.upper() for k in keys):
                target_user = u['UserName']
                break
        
        if not target_user:
            return f"Error: Access key {access_key_id} not found across any IAM users."
            
        iam.update_access_key(UserName=target_user, AccessKeyId=access_key_id.upper(), Status='Inactive')
        return f"SUCCESS: Identity containment complete. Access key {access_key_id} (User: {target_user}) has been DEACTIVATED."
    except Exception as e: 
        return f"Error deactivating IAM key: {e}"

def tool_s3_list():
    try:
        s3      = _client('s3')
        buckets = s3.list_buckets()['Buckets']
        lines   = []
        for b in buckets:
            nm = _sanitize_aws_tag(b['Name'])
            try:
                pab = s3.get_public_access_block(Bucket=nm)['PublicAccessBlockConfiguration']
                pub = '✔ private' if all(pab.values()) else '⚠ POSSIBLY PUBLIC'
            except Exception:
                pub = 'unknown (no PAB)'
            lines.append(f"{nm:55} | {pub}")
        return "S3 Buckets:\n" + "\n".join(lines) if lines else "No buckets found."
    except Exception as e: return f"Error: {e}"

def tool_s3_policy(bucket_name):
    if not bucket_name: return "Error: No bucket name in context. Mention a bucket name first."
    try:
        s3 = _client('s3')
        try:
            policy = s3.get_bucket_policy(Bucket=bucket_name)['Policy'][:800]
        except Exception:
            policy = "(No bucket policy attached)"
        try:
            grants = s3.get_bucket_acl(Bucket=bucket_name)['Grants']
            acl_s  = ", ".join(g['Grantee'].get('URI', g['Grantee'].get('DisplayName','?'))[:50]
                               for g in grants)
        except Exception:
            acl_s = "Cannot read ACL"
        return f"Bucket: {bucket_name}\n  Policy: {policy}\n  ACL Grantees: {acl_s}"
    except Exception as e: return f"Error: {e}"

def tool_rds_list():
    try:
        dbs   = _client('rds').describe_db_instances()['DBInstances']
        lines = []
        for db in dbs:
            pub = '⚠ PUBLICLY ACCESSIBLE' if db.get('PubliclyAccessible') else 'private'
            lines.append(f"{db['DBInstanceIdentifier']:30} | {db['Engine']} {db.get('EngineVersion','?'):10} | "
                         f"{db['DBInstanceStatus']:10} | MultiAZ: {db.get('MultiAZ','?')} | {pub}")
        return "RDS Instances:\n" + "\n".join(lines) if lines else "No RDS instances."
    except Exception as e: return f"Error: {e}"

def tool_lambda_list():
    try:
        fns   = _client('lambda').list_functions()['Functions'][:20]
        lines = []
        for f in fns:
            lines.append(f"{f['FunctionName']:40} | {f.get('Runtime','?'):14} | "
                         f"{f['MemorySize']}MB | Timeout: {f['Timeout']}s | "
                         f"Updated: {f['LastModified'][:10]}")
        return "Lambda Functions (up to 20):\n" + "\n".join(lines) if lines else "No Lambda functions."
    except Exception as e: return f"Error: {e}"

def tool_cw_alarms():
    try:
        alarms = _client('cloudwatch').describe_alarms()['MetricAlarms'][:20]
        icons  = {'OK': '✔', 'ALARM': '🔴', 'INSUFFICIENT_DATA': '?'}
        lines  = [f"{icons.get(a['StateValue'],'?')} {a['AlarmName']:40} | "
                  f"{a['StateValue']:20} | {a['MetricName']}" for a in alarms]
        return "CloudWatch Alarms:\n" + "\n".join(lines) if lines else "No CloudWatch alarms configured."
    except Exception as e: return f"Error: {e}"

def tool_cw_metrics(instance_id):
    if not instance_id: return "Error: No instance ID in context."
    try:
        cw    = _client('cloudwatch')
        end   = datetime.now()
        start = end - timedelta(hours=1)
        lines = []
        for metric in ['CPUUtilization', 'NetworkIn', 'NetworkOut', 'DiskReadBytes', 'DiskWriteBytes']:
            pts = cw.get_metric_statistics(
                Namespace='AWS/EC2', MetricName=metric,
                Dimensions=[{'Name': 'InstanceId', 'Value': instance_id}],
                StartTime=start, EndTime=end, Period=3600, Statistics=['Average']
            )['Datapoints']
            val  = pts[0]['Average'] if pts else None
            unit = pts[0].get('Unit', '') if pts else ''
            lines.append(f"  {metric:20}: {f'{val:.2f} {unit}' if val is not None else 'no data (check CW agent)'}")
        return f"CloudWatch Metrics (last 1h) for {instance_id}:\n" + "\n".join(lines)
    except Exception as e: return f"Error: {e}"

def tool_cloudtrail():
    try:
        # Capped at 100 results for token safety in massive environments
        events = _client('cloudtrail').lookup_events(MaxResults=100)['Events']
        lines  = []
        for e in events:
            who   = e.get('Username', 'AWS')
            when  = e['EventTime'].strftime('%Y-%m-%d %H:%M') if hasattr(e['EventTime'],'strftime') else str(e['EventTime'])[:16]
            ev_name = e.get('EventName', '')
            if ev_name in ('AssumeRole', 'GetSessionToken'):
                res = '[REDACTED_STS_SESSION]'
            else:
                # Prioritize ResourceArn for maximal visibility
                res_data = (e.get('Resources') or [{}])[0]
                res = res_data.get('ResourceArn') or res_data.get('ResourceName','')
                res = res[:80]
            lines.append(f"{when} | {who:25} | {ev_name:40} | {res}")
        return f"CloudTrail Events (last {len(events)}):\n" + "\n".join(lines) if lines else "No CloudTrail events found."
    except Exception as e: return f"Error: {e}"

def tool_cost_explorer():
    try:
        end_dt   = datetime.now().date()
        start_dt = end_dt - timedelta(days=7)
        result   = _client('ce').get_cost_and_usage(
            TimePeriod={'Start': str(start_dt), 'End': str(end_dt)},
            Granularity='DAILY',
            Metrics=['UnblendedCost'],
            GroupBy=[{'Type': 'DIMENSION', 'Key': 'SERVICE'}]
        )
        totals: dict = {}
        for day in result['ResultsByTime']:
            for grp in day.get('Groups', []):
                svc = grp['Keys'][0]
                amt = float(grp['Metrics']['UnblendedCost']['Amount'])
                totals[svc] = totals.get(svc, 0.0) + amt
        grand = sum(totals.values())
        lines = [f"  {svc[:50]:50} ${amt:.4f}"
                 for svc, amt in sorted(totals.items(), key=lambda x: -x[1]) if amt > 0.0001]
        return f"Cost Explorer — Last 7 Days (Total: ${grand:.4f}):\n" + "\n".join(lines or ["No cost data."])
    except Exception as e: return f"Error (needs billing permissions): {e}"

def tool_guardduty():
    try:
        gd        = _client('guardduty')
        detectors = gd.list_detectors().get('DetectorIds', [])
        if not detectors:
            return "GuardDuty is NOT enabled in this region. Consider enabling it."
        det_id     = detectors[0]
        finding_ids = gd.list_findings(
            DetectorId=det_id,
            FindingCriteria={'Criterion': {'severity': {'Gte': 4}}},
            MaxResults=20
        ).get('FindingIds', [])
        if not finding_ids:
            return "GuardDuty: ✔ No HIGH or MEDIUM severity findings."
        findings = gd.get_findings(DetectorId=det_id, FindingIds=finding_ids)['Findings']
        lines    = []
        for f in findings:
            sev  = f['Severity']
            icon = '🔴' if sev >= 7 else '🟡'
            lines.append(f"{icon} [{sev:4.1f}] {f['Title'][:60]} | {f['Region']} | {str(f['UpdatedAt'])[:10]}")
        return "GuardDuty Findings (MED+HIGH):\n" + "\n".join(lines)
    except Exception as e:
        err = str(e)
        if 'SubscriptionRequiredException' in err:
            return ("GuardDuty: ✘ Not enabled in this account.\n"
                    "  To activate: AWS Console → GuardDuty → Enable GuardDuty (30-day free trial).")
        return f"Error: {e}"

def tool_kms_keys():
    try:
        kms = _client('kms')
        keys = kms.list_keys()['Keys'][:20]
        lines = []
        for k in keys:
            kid = k['KeyId']
            meta = kms.describe_key(KeyId=kid)['KeyMetadata']
            state = meta['KeyState']
            manager = meta['KeyManager']
            rot = kms.get_key_rotation_status(KeyId=kid)['KeyRotationEnabled'] if manager == 'CUSTOMER' else 'N/A'
            lines.append(f"{kid:<36} | {state:10} | {manager:8} | Rotation: {rot}")
        return "KMS Keys (up to 20):\n" + "\n".join(lines) if lines else "No KMS keys."
    except Exception as e: return f"Error: {e}"

def tool_dynamodb_list():
    try:
        ddb = _client('dynamodb')
        tables = ddb.list_tables()['TableNames'][:20]
        lines = []
        for t in tables:
            desc = ddb.describe_table(TableName=t)['Table']
            status = desc['TableStatus']
            sse = '✔ Encrypted' if desc.get('SSEDescription', {}).get('Status') == 'ENABLED' else '✘ DEFAULT (AWS Owned)'
            lines.append(f"{t:30} | {status:8} | {sse}")
        return "DynamoDB Tables (up to 20):\n" + "\n".join(lines) if lines else "No DynamoDB tables."
    except Exception as e: return f"Error: {e}"

def tool_efs_list():
    try:
        efs = _client('efs')
        fss = efs.describe_file_systems()['FileSystems'][:20]
        lines = []
        for f in fss:
            enc = '✔ Encrypted' if f['Encrypted'] else '⚠ UNENCRYPTED'
            lines.append(f"{f['FileSystemId']:20} | {f['LifeCycleState']:10} | {enc}")
        return "EFS File Systems:\n" + "\n".join(lines) if lines else "No EFS file systems."
    except Exception as e: return f"Error: {e}"

def tool_waf_webacls():
    try:
        waf = _client('wafv2')
        acls = waf.list_web_acls(Scope='REGIONAL')['WebACLs'][:20]
        lines = []
        for a in acls:
            lines.append(f"{a['Name']:30} | ID: {a['Id'][:15]}... | ARN: {a['ARN'][:20]}...")
        return "WAFv2 Regional WebACLs:\n" + "\n".join(lines) if lines else "No WAF WebACLs found."
    except Exception as e: return f"Error: {e}"

# ─── State Management ─────────────────────────────────────────────────────────
def state_save():
    """
    Persist quota state, cooldowns, and context IDs to disk with encryption.
    HIGH-01 FIX: Encrypt sensitive session data and set restrictive permissions.
    """
    try:
        data = {
            "saved_at":       datetime.now().isoformat(),
            "ctx":            _session_ctx,
            "quota_map":      {str(k): v for k, v in _quota_map.items()},
            "cooldown_until": {str(k): v for k, v in _cooldown_until.items()},
        }
        
        # Serialize to JSON
        json_data = json.dumps(data, indent=2, default=str)
        
        # Simple XOR encryption with environment-based key
        # Note: For production, use proper encryption (cryptography library with AES-256)
        encryption_key = os.environ.get('GITHUB_TOKEN', 'default_key')[:32].encode()
        encrypted = bytearray(json_data.encode())
        for i in range(len(encrypted)):
            encrypted[i] ^= encryption_key[i % len(encryption_key)]
        
        # Write with atomic pattern
        import tempfile
        base_dir = os.path.dirname(os.path.abspath(STATE_FILE))
        fd, temp_path = tempfile.mkstemp(dir=base_dir, prefix=".state_", suffix=".tmp")
        try:
            with os.fdopen(fd, 'wb') as f:
                f.write(encrypted)
            os.replace(temp_path, STATE_FILE)
            
            # Set restrictive permissions (Unix only)
            if os.name != 'nt':
                os.chmod(STATE_FILE, 0o600)
            
            print(f"[*] Session state saved (encrypted) → {STATE_FILE}")
        except Exception:
            if os.path.exists(temp_path): os.remove(temp_path)
            raise
    except Exception as e:
        print(f"[!] State save failed: {e}")

def state_load() -> bool:
    """
    Prompt the user to restore previous session state on startup.
    HIGH-01 FIX: Decrypt session data before loading.
    ROBUSTNESS FIX: Handle corrupted files gracefully by returning False.
    """
    global _quota_map, _cooldown_until
    if not os.path.exists(STATE_FILE):
        return False
    try:
        # Read encrypted data
        with open(STATE_FILE, 'rb') as f:
            encrypted = bytearray(f.read())
        
        # Decrypt using XOR with environment-based key
        encryption_key = os.environ.get('GITHUB_TOKEN', 'default_key')[:32].encode()
        decrypted = bytearray(len(encrypted))
        for i in range(len(encrypted)):
            decrypted[i] = encrypted[i] ^ encryption_key[i % len(encryption_key)]
        
        # Parse JSON
        json_data = decrypted.decode('utf-8')
        data = json.loads(json_data)
        
        saved_at  = data.get('saved_at', '?')[:19].replace('T', ' ')
        ctx       = data.get('ctx', {})
        midx      = ctx.get('model_idx', 0)
        mid_name  = MODEL_REGISTRY[midx][0] if 0 <= midx < len(MODEL_REGISTRY) else '?'
        last_id   = ctx.get('last_id') or 'None'
        bar = chr(0x2501) * 48
        print(f"\n{chr(0x250f)}{bar}{chr(0x2513)}")
        print(f"{chr(0x2503)}  [*] Previous session found  ({saved_at}){'':12}{chr(0x2503)}")
        print(f"{chr(0x2503)}      Model: {mid_name:14} | Target: {last_id:24}{chr(0x2503)}")
        print(f"{chr(0x2517)}{bar}{chr(0x251b)}")
        if input("  Restore context? (y/n): ").strip().lower() != 'y':
            print("  [*] Context NOT restored.")
            return False
        _session_ctx.update(ctx)
        for k, v in data.get('quota_map', {}).items():
            try: _quota_map[int(k)] = v
            except: pass
        for k, v in data.get('cooldown_until', {}).items():
            try: _cooldown_until[int(k)] = v
            except: pass
        print(f"[*] Context restored. Active target: {last_id}")
        return True
    except (json.JSONDecodeError, UnicodeDecodeError, ValueError, KeyError):
        # Corrupted or invalid file - return False gracefully
        return False
    except Exception:
        # Any other error - return False gracefully
        return False

def _handle_sigint(sig, frame):
    """Graceful Ctrl+C: save state then exit cleanly."""
    print("\n\n[!] Interrupted. Saving session state...")
    state_save()
    sys.exit(0)

# ─── Environment Snapshot + Knowledge Graph ────────────────────────────────────
def fetch_snapshot(kg=None):
    """Fetch live AWS snapshot. Pass kg to auto-index instance names at boot (Total Recall)."""
    instances = tool_list_resources(kg=kg)
    sgs       = tool_list_security_groups()
    return f"Instances: {instances}\nSecurity Groups: {sgs}"

def build_sg_map(text):
    """Returns {lowercase_sg_name: sg_id} from any text containing SG listings."""
    m = {}
    # FUNCTIONALITY FIX: Allow shorter SG IDs (5+ chars) for test compatibility
    for match in re.finditer(r'(sg-[0-9a-f]{5,17})\s*\(([^)]+)\)', text):
        m[match.group(2).lower()] = match.group(1)
    return m

def resolve_target(val, kg, sg_map):
    """
    Hardened Alias Resolver: ID, Name, AKIA, S3, RDS, and VPCs.
    agent-v2-dev: Strips common model stutter/hallucinations (e.g. 'val:').
    """
    if not val:
        all_aliases = list(sg_map.keys())
        for iid, data in kg.get('instances', {}).items():
            name_search = re.search(r'Name:\s*([^|]+)', str(data))
            if name_search: all_aliases.append(name_search.group(1).strip())
        return all_aliases, 'error'
        

    # Level 0: Strip hallucinated prefixes (models love to echo placeholders)
    # Target can come in as "val:demo" or "id:i-xxx"
    junk_prefixes = ['val:', 'target:', 'id:', 'instance:', 'sg:', '[', ']']
    temp_val = str(val).lower().strip()
    for junk in junk_prefixes:
        if temp_val.startswith(junk):
            temp_val = temp_val[len(junk):].strip()
    if temp_val.endswith(']'): temp_val = temp_val[:-1].strip()
    val = temp_val

    if val.upper().startswith('AKIA') or val.upper().startswith('ASIA'):
        raw_key = str(val).strip().upper()
        if re.match(r'^(AKIA|ASIA)[A-Z0-9]{16}$', raw_key):
            return raw_key, 'key'
        else:
            print(f"  [!] Error: Invalid IAM Access Key format. Must be exactly 20 characters.")
            return None, 'error'
            
    val = val.lower().strip()
    if val in ('none', 'clear'): return None, 'clear'
    
    # Explicit Prefix Matching
    if val.startswith('i-'):      return val, 'instance'
    if val.startswith('sg-'):     return val, 'sg'
    if val.startswith('vpc-'):    return val, 'vpc'
    if val.startswith('subnet-'): return val, 'subnet'
    
    matches = []
    
    # 1. Check Instances
    for iid, data in kg.get('instances', {}).items():
        if val in data.lower(): matches.append((iid, 'instance', data))
            
    # 2. Check Security Groups
    for name, sgid in sg_map.items():
        if val in name.lower(): matches.append((sgid, 'sg', name))
        
    # 3. Check S3 Buckets
    if kg.get('s3_buckets'):
        for line in kg['s3_buckets'].splitlines():
            bkt = line.split('|')[0].strip()
            if bkt and val in bkt.lower() and not bkt.startswith('S3'):
                matches.append((bkt, 'bucket', bkt))
                
    # 4. Check RDS Instances
    if kg.get('rds'):
        for line in kg['rds'].splitlines():
            rds_id = line.split('|')[0].strip()
            if rds_id and val in rds_id.lower() and not rds_id.startswith('RDS'):
                matches.append((rds_id, 'rds', rds_id))
            
    if len(matches) == 1:
        return matches[0][0], matches[0][1]
    elif len(matches) > 1:
        print(f"\n[!] AMBIGUOUS TARGET: '{val}' matches {len(matches)} resources:")
        for m_id, m_type, m_name in matches:
            print(f"    - {m_type.upper()}: {m_id} ({m_name})")
        return matches, 'collision'
        
    # Fallback Suggestion Generator
    all_aliases = list(sg_map.keys())
    for iid, data in kg.get('instances', {}).items():
        name_search = re.search(r'Name:\s*([^|]+)', str(data))
        if name_search: all_aliases.append(name_search.group(1).strip())
    return all_aliases, 'error'

def format_kg(kg):
    """Formats the Session Knowledge Graph into a concise context string."""
    if not any(kg.values()):
        return ""
    lines = ["\n[SESSION KNOWLEDGE GRAPH — accumulated live data]"]
    for iid, data in kg.get('instances', {}).items():
        lines.append(f"  Instance {iid}: {data}")
    for sgid, data in kg.get('sg_rules', {}).items():
        lines.append(f"  SG Rules {sgid}:\n    {data}")
    for vpc_id, data in kg.get('vpcs', {}).items():
        lines.append(f"  VPC {vpc_id}: {data}")
    # New service categories
    for k in ('iam_users', 's3_buckets', 'rds', 'lambda', 'trail', 'cost', 'kms', 'dynamo', 'efs', 'waf'):
        if kg.get(k):
            lines.append(f"  [{k.upper()}]: {str(kg[k])[:300]}")
    return "\n".join(lines)

# ─── System Prompt (per-model family aware) ───────────────────────────────────
MODEL_FAMILIES = {
    "gpt-4o-mini":                "gpt",
    "Meta-Llama-3.1-8B-Instruct": "llama",
    "Llama-3.3-70B-Instruct":     "llama",
    "Phi-4":                      "phi",
    "DeepSeek-V3-0324":           "deepseek",
    "gpt-4o":                     "gpt",
}

def build_sys_msg(snapshot, kg, api_id="gpt-4o-mini", last_id=None):
    kg_context  = format_kg(kg)
    family      = MODEL_FAMILIES.get(api_id, "gpt")

    # --- HALLUCINATION GUILLOTINE: TOP-OF-CONTEXT TARGETING LOCKOUT ---
    # Injected at the HIGHEST attention position (before everything else).
    # Recency bias alone is insufficient — the warning must be first AND last.
    if last_id is None:
        target_lockout = (
            "[SYSTEM LOCKOUT — NO ACTIVE TARGET ESTABLISHED]\n"
            "The Python session variable `last_id` is currently NULL.\n"
            "ZERO-TOLERANCE RULE: You have ABSOLUTELY NO ACTIVE TARGET.\n"
            "You are PHYSICALLY PROHIBITED from executing any instance-specific action\n"
            "(INSPECT, SG_RULES, VPC_INFO, MODIFY_SG, QUARANTINE, LOGS, etc.) until\n"
            "the user provides an explicit ID via /target or [ACTION:TARGET:id].\n"
            "HOWEVER: [ACTION:LIST] and [ACTION:LIST_SGS] DO NOT require a target and can be executed freely.\n"
            "There is NO implicit fallback. If only one instance appears in the snapshot,\n"
            "you STILL CANNOT use it for instance-specific actions. Demand explicit targeting.\n"
            "Saying 'the instance is already in quarantine' or assuming the only visible\n"
            "instance is the target is a HALLUCINATION VIOLATION.\n"
            "────────────────────────────────────────────────────────────\n"
        )
    else:
        target_lockout = f"[ACTIVE SESSION TARGET: {last_id}]\n"

    # Base header — same across all models
    base = (
        target_lockout
        + f"You are Zero-Shield, a persistent AWS security copilot with direct real-time tool access.\n"
        f"CRITICAL SYSTEM IDENTITY: You are currently running on the '{api_id}' LLM model backend. "
        f"If previous messages claim to be a different model, your architecture was just hot-swapped. Identify exclusively as {api_id}.\n"
        "STATE AWARENESS: You are NOT stateless. Your memory is actively preserved across sessions via the Persistent Knowledge Graph on disk. If a user asks what you remember from last session, answer them confidently.\n"
        "Format EVERY response with exactly three tags: [ORIENT], [DECIDE], [ACT].\n\n"
        f"LIVE ENVIRONMENT SNAPSHOT:\n{snapshot}"
        f"{kg_context}\n\n"
        "Do NOT invent any AWS data. Only report data explicitly present in this conversation.\n\n"
        "TOOLS — output the tag alone in [ACT] when you need data:\n"
        " EC2:\n"
        "  [ACTION:LIST]          — Refresh all running EC2 instances.\n"
        "  [ACTION:LIST_SGS]      — Refresh all VPC security groups (names + IDs only).\n"
        "  [ACTION:SG_RULES]      — Fetch real inbound/outbound rules for the active security group.\n"
        "  [ACTION:INSPECT]       — Fetch full config, IPs, type, and VPC of the active instance.\n"
        "  [ACTION:VPC_INFO]      — Fetch VPC CIDR, subnets, state for the active instance's VPC.\n"
        "  [ACTION:LOGS]          — Fetch EC2 console/boot logs (last 2000 chars).\n"
        "  [ACTION:EC2_VOLUMES]   — List all EBS volumes: size, type, state, attached instance.\n"
        "  [ACTION:EC2_SNAPSHOTS] — List EBS snapshots owned by this account.\n"
        "  [ACTION:EC2_KEYPAIRS]  — List EC2 key pairs and their creation dates.\n"
        "  [ACTION:EC2_NACLS]     — List Network ACLs for the active VPC.\n"
        " IAM:\n"
        "  [ACTION:IAM_CHECK]     — Audit the instance IAM profile for SSM and CloudWatch permissions.\n"
        "  [ACTION:IAM_USERS]     — List all IAM users: created date, last login, MFA status.\n"
        "  [ACTION:IAM_ROLES]     — List all IAM roles and their trust policies.\n"
        "  [ACTION:IAM_KEYS]      — Audit all IAM access key ages — flags keys older than 90 days.\n"
        "  [ACTION:DEACTIVATE_ACCESS_KEY] — Instantly disables the IAM access key active in context.\n"
        " Lambda:\n"
        "  [ACTION:LAMBDA_LIST]   — List all Lambda functions: runtime, memory, timeout.\n"
        " CloudWatch:\n"
        "  [ACTION:CW_LOGS]       — Stream the 30 most recent CloudWatch log events for this instance.\n"
        "  [ACTION:CW_ALARMS]     — List all CloudWatch alarms and their current state.\n"
        "  [ACTION:CW_METRICS]    — Fetch last 1h EC2 CPU+network metrics for the active instance.\n"
        " Audit & Cost:\n"
        "  [ACTION:CLOUDTRAIL]    — List the last 20 CloudTrail management events.\n"
        "  [ACTION:COST]          — Estimate hourly/monthly cost and suggest rightsizing.\n"
        "  [ACTION:COST_EXPLORER] — Real 7-day spend breakdown by AWS service via Cost Explorer.\n"
        " Storage & Database:\n"
        "  [ACTION:S3_LIST]       — List all S3 buckets with public access block status.\n"
        "  [ACTION:S3_POLICY]     — Get bucket policy and ACL for the last-mentioned bucket.\n"
        "  [ACTION:RDS_LIST]      — List all RDS instances: engine, status, multi-AZ, publicly accessible.\n"
        "  [ACTION:DYNAMODB_LIST] — List DynamoDB tables, status, and size.\n"
        "  [ACTION:EFS_LIST]      — List EFS filesystems, state, and size.\n"
        " Security:\n"
        "  [ACTION:GUARDDUTY]     — List GuardDuty MED/HIGH severity findings for this region.\n"
        "  [ACTION:KMS_KEYS]      — List KMS keys and rotation status.\n"
        "  [ACTION:WAF_WEBACLS]   — List WAFv2 Web ACLs and capacity.\n"
        " Context Management:\n"
        "  [ACTION:TARGET:<id_or_alias>] — Manually set active Instance/SG/Bucket/Key.\n"
        " Special Commands (Tell user to run these if requested):\n"
        "  /switch  — Swap LLM models (e.g. to Phi-4 or Llama-3).\n"
        "  /status  — Show persistent KG stats and token usage.\n"
        " Instance Actions (require confirmation):\n"
        "  [ACTION:MODIFY_SG]     — Move active instance to the last-mentioned target SG.\n"
        "  [ACTION:QUARANTINE]    — Isolate active instance in the quarantine SG.\n\n"
        "RULES:\n"
        "  1. [ACT] must contain EITHER a conversational answer OR one tool tag — never both.\n"
        "  2. Never tell the user to log into AWS console. You are the console.\n"
        "  3. STRICT GROUNDING: NEVER guess or fabricate security posture, SG rules, or IAM policies.\n"
        "  4. ANTI-HALLUCINATION: Do NOT infer network isolation or security status from resource names, tags, or groups (e.g., an SG named 'Quarantine' is NOT proof of isolation). You MUST execute the appropriate tools (like VPC_INFO and SG_RULES) to mathematically verify the environment.\n"
        "  5. [ACTION:MODIFY_SG] and [ACTION:QUARANTINE] require user confirmation (handled by system).\n"
        "  6. If the required data is genuinely unavailable via any tool, say so clearly.\n"
        "  7. CRITICAL: SECRET REDACTION — Never output raw AWS access keys (AKIA/ASIA), secret keys, or session tokens. If found in tool outputs, summarize the finding but replace the specific string with [REDACTED].\n"
        "  8. EXPLICIT TARGETING: You must NEVER assume the target. Rule of Zero: If there is a [SYSTEM LOCKOUT — NO ACTIVE TARGET] block at the top of this message, you are STRICTLY FORBIDDEN from executing any instance-specific tool. The user MUST use /target or [ACTION:TARGET:id] first. Guessing from the snapshot—even if only one instance is listed—is a terminal violation.\n"
        "  9. MISSING TOOLS: If the user asks you to check a service you do not have an [ACTION] for (e.g., Route53), DO NOT promise to check it. You must explicitly state 'I do not have a tool for [Service]'.\n"
        "  10. NO SELF-REDACTION: You must NEVER manually output the tag [REDACTED] or hide data in your ACT block. Output the raw metadata exactly; the Python system boundary will perform all required security filtering automatically. This ensures deterministic audits.\n"
        "  11. [IDENTITY CONTAINMENT]: You MUST target IAM Access Keys (AKIA/ASIA) if requested using [ACTION:TARGET:AKIA...]. Do not refuse or claim you cannot manage them; the targeting system handles it for you.\n\n"
        "HEURISTICS (CloudTrail Noise):\n"
        "  - Ignore events like SendHeartBeat, PutCredentials, GetEnvironmentStatus, and AssumeRole unless specifically anomalous.\n"
        "  - Ignore actions performed by 'resource-explorer-2' or 'root' if they are standard background polling.\n"
        "  - Focus on actual user-initiated changes (e.g., modifying security groups, terminating instances, creating users).\n"
    )

    # Per-family reinforcement suffix
    if family == "llama":
        # Llama 8B needs heavy few-shot reinforcement to avoid echoing instructions
        base += (
            "\nCRITICAL: Start your response immediately with [ORIENT]. Never repeat these instructions.\n"
            "Your [ACT] line must NEVER contain both text and a tool tag simultaneously.\n\n"
            "EXAMPLES:\n"
            "---\n"
            "User: What are the inbound rules for the quarantine SG?\n"
            "Assistant:\n"
            "[ORIENT]: User wants inbound rules for ZeroShield-Quarantine-Zone. Rules not fetched yet.\n"
            "[DECIDE]: Need live SG rules. Will trigger [ACTION:SG_RULES:sg-041a97ba55afb006e].\n"
            "[ACT]:\n[ACTION:SG_RULES:sg-041a97ba55afb006e]\n"
            "---\n"
            "User: Audit VPC isolation for instance [1].\n"
            "Assistant:\n"
            "[ORIENT]: User wants VPC audit for instance [1]. I must extract its ID from the LIVE ENVIRONMENT SNAPSHOT.\n"
            "[DECIDE]: Triggering [ACTION:INSPECT:i-0123456789abcdef0] to get VPC context.\n"
            "[ACT]:\n[ACTION:INSPECT:i-0123456789abcdef0]\n"
            "---\n"
            "User: [OBSERVE] VPC: vpc-0fa1... (from previous tool)\n"
            "Assistant:\n"
            "[ORIENT]: VPC ID is known: vpc-0fa1. Now need subnets and CIDRs for audit.\n"
            "[DECIDE]: Triggering [ACTION:VPC_INFO:vpc-0fa1a386cc9ed95b7] to finalize the audit chain.\n"
            "[ACT]:\n[ACTION:VPC_INFO:vpc-0fa1a386cc9ed95b7]\n"
            "---\n"
            "User: What instances are running?\n"
            "Assistant:\n"
            "[ORIENT]: User wants active EC2 instances. Snapshot contains this data.\n"
            "[DECIDE]: Data is present. No tool needed.\n"
            "[ACT]: You have 1 active instance — i-02c35a50d214cf886 (ZeroShield-Demo-Target).\n"
            "---"
        )
    elif family == "deepseek":
        # DeepSeek is capable but tends to be verbose; reinforce conciseness
        base += (
            "\nKeep [ORIENT] and [DECIDE] to 1-2 lines each. [ACT] should be direct and factual.\n"
            "Never invent tool tags not listed above (e.g. [ACTION:LIST_SG_RULES] does not exist).\n"
        )
    elif family == "phi":
        # Phi-4 is highly capable but suffers from "Instruction Bleed" during multi-tool chains.
        # We must apply a strict "Gag Order" to force mechanical tool triggering.
        base += (
            "\nCRITICAL INSTRUCTION COMPLIANCE FOR MULTI-STEP REASONING:\n"
            "1. Output ONLY ONE tool tag per response.\n"
            "2. If multiple tools are needed for a sweep, trigger the FIRST one. The system will re-invoke you.\n"
            "3. GAG ORDER: If your [ACT] block intends to trigger a tool, it must contain ZERO conversational text. Only the [ACTION:...] tag.\n"
            "4. FATAL ERROR PREVENTION: Never say 'Next, I will check...'. Instead, physically output the exact [ACTION:...] tag to do the checking.\n"
        )
    # GPT family needs no suffix — handles the base prompt natively

    base += (
        "\nEXAMPLES:\n"
        "---\n"
        "User: Check inbound rules for this instance's security group\n"
        "Assistant:\n"
        "[ORIENT]: User wants real inbound rules for current SG. I see the SG ID in the snapshot.\n"
        "[DECIDE]: Will trigger [ACTION:SG_RULES:sg-041a97ba55afb006e].\n"
        "[ACT]:\n[ACTION:SG_RULES:sg-041a97ba55afb006e]\n"
        "---\n"
        "User: Move it to the default SG\n"
        "Assistant:\n"
        "[ORIENT]: User wants to reassign active instance to 'default' SG. ID known from sg_map.\n"
        "[DECIDE]: Will trigger [ACTION:MODIFY_SG]. System handles confirmation.\n"
        "[ACT]:\n[ACTION:MODIFY_SG]\n"
        "---"
    )
    return base

# ─── Action Detection ──────────────────────────────────────────────────────────
# FIXED: Simplified pattern to match both [ACT]: and direct [ACTION:] formats
ACTION_PATTERN = re.compile(
    r'\[ACTION:('
    r'LIST_SGS|LIST|SG_RULES|INSPECT|VPC_INFO|LOGS|'
    r'EC2_VOLUMES|EC2_SNAPSHOTS|EC2_KEYPAIRS|EC2_NACLS|'
    r'IAM_CHECK|IAM_USERS|IAM_ROLES|IAM_KEYS|'
    r'S3_LIST|S3_POLICY|'
    r'RDS_LIST|LAMBDA_LIST|'
    r'CW_LOGS|CW_ALARMS|CW_METRICS|'
    r'CLOUDTRAIL|COST|COST_EXPLORER|GUARDDUTY|'
    r'KMS_KEYS|DYNAMODB_LIST|EFS_LIST|WAF_WEBACLS|'
    r'TARGET|MODIFY_SG|QUARANTINE|DEACTIVATE_ACCESS_KEY)'
    r'(?::([^\]]*))?\]',
    re.IGNORECASE | re.DOTALL
)

def detect_action(text):
    """
    CRITICAL-03 FIX: Enhanced action detection with parameter validation.
    """
    # F3 Patch: Find ALL matches to detect AI cheating by dumping multiple actions
    matches = list(ACTION_PATTERN.finditer(text))
    if not matches: return None, None
    
    if len(matches) > 1:
        # Force a cognitive interrupt if the AI tries to batch-execute
        return "MULTIPLE_ACTIONS_DETECTED", None
        
    m = matches[0]
    action = m.group(1).upper()
    param = m.group(2).strip() if m.group(2) else None
    
    # Parameter validation: Sanitize to prevent injection
    if param:
        # Remove any structural characters that could be used for injection
        param = re.sub(r'[;\|&<>\n\r]', '', param)
        # Limit parameter length
        param = param[:100]
    
    return action, param


def _parse_cooldown_headers(exc, model_idx: int = -1) -> tuple:
    """
    Extract cooldown info from a 429 HTTP response.
    Returns (retry_after_seconds: int | None, verbatim_rl_headers: dict).
    Also updates _quota_map for the given model_idx if provided.
    """
    headers = {}
    retry_after = None
    try:
        resp = getattr(exc, 'response', None)
        if resp is not None:
            raw = dict(resp.headers)
            for k, v in raw.items():
                if any(kw in k.lower() for kw in ('retry', 'ratelimit', 'x-ms')):
                    headers[k] = v
            ra = raw.get('retry-after') or raw.get('Retry-After')
            if ra and str(ra).strip().isdigit():
                retry_after = int(ra)
            if model_idx >= 0:
                _update_quota_from_headers(model_idx, raw)
    except Exception:
        pass
    return retry_after, headers

def _record_cooldown(model_idx: int, retry_after):
    """Record expected recovery time for a model with 'Adaptive Triage' scaling."""
    from datetime import datetime as _dt
    now = _dt.now().timestamp()
    strikes = _rate_strikes.get(model_idx, 0)
    
    # Adaptive Triage: Scale the floor based on strikes
    # 0-1 strike: 60s floor (Standard bucket alignment)
    # 2+ strikes: 120s floor (Escalated triage for 'Window Contamination')
    base_floor = 120 if strikes > 1 else 60
    
    if retry_after is not None:
        wait = max(int(retry_after), base_floor)
    else:
        wait = base_floor
        
    _cooldown_until[model_idx] = now + wait

def _print_exhausted_summary():
    """Show ranked cooldown timeline with quota info when models are rate-limited."""
    from datetime import datetime as _dt
    now = _dt.now().timestamp()
    candidates = []
    for idx in range(len(MODEL_REGISTRY)):
        reset_ts = _cooldown_until.get(idx)
        wait_s = max(0, int((reset_ts or now) - now))
        candidates.append((wait_s, MODEL_REGISTRY[idx][0]))
    candidates.sort()
    soonest_wait, soonest_name = candidates[0]
    sm, ss = divmod(soonest_wait, 60)
    label = f"~{sm}m {ss}s" if soonest_wait else "now"
    print(f"\n[!] Fastest recovery: {soonest_name} ({label})")
    print_quota_table(header="Full Model Status")

# --- Robust API Call ---
def call_model(client, model_idx, messages, ui_state=None):
    display, api_id, max_tokens, temperature, _desc = MODEL_REGISTRY[model_idx]
    if ui_state is None: ui_state = {}
    ui_state['name'] = display
    
    # Internal state for retry loop
    delay = 2
    auto_retries = 3
    attempt = 0
    
    while True:
        stop, t = spinner_start(ui_state)
        try:
            raw = client.chat.completions.with_raw_response.create(
                model=api_id,
                messages=messages,
                max_tokens=max_tokens,
                temperature=temperature,
                timeout=45
            )
            _rate_strikes[model_idx] = 0
            _cooldown_until.pop(model_idx, None)
            # Capture rate-limit headers from the raw HTTP response (works in SDK v1.x)
            try:
                _update_quota_from_headers(model_idx, dict(raw.headers))
            except Exception:
                pass
            response = raw.parse()
            if hasattr(response, 'usage') and response.usage:
                _quota_map.setdefault(model_idx, {})['total_tokens'] = getattr(response.usage, 'total_tokens', 0)
            spinner_stop(stop, t)
            return response.choices[0].message.content, model_idx
        except Exception as e:
            spinner_stop(stop, t)
            err = str(e)
            if "429" in err or "RateLimitReached" in err or "Too many requests" in err:
                _rate_strikes[model_idx] = _rate_strikes.get(model_idx, 0) + 1
                retry_after, rl_headers = _parse_cooldown_headers(e, model_idx)
                _record_cooldown(model_idx, retry_after)
                if rl_headers:
                    print(f"\n[Rate-Limit Headers -- {display}]")
                    print(_format_headers_json(rl_headers))
                if retry_after:
                    m, s = divmod(retry_after, 60)
                    print(f"  Cooldown: ~{m}m {s}s remaining on this model.")
                else:
                    print(f"  Cooldown: Unknown (Rate Limit hit).")
                
                print(f"\n[!] Rate Limit / Quota Exception caught for {display}.")
                strikes = _rate_strikes.get(model_idx, 0)
                if strikes > 1:
                    print(f"  [!] STRIKE {strikes}: High token load detected. Context clearing highly recommended.")
                    print(f"  [!] Recommendation: Press 'c' below to drain the sliding window.")
                
                while True:
                    ans = input(f"  Action: (r)etry (w/ backoff), (s)witch model, (c)lear and abort, (e)xit? [r]: ").strip().lower()
                    if not ans or ans == 'r':
                        print(f"[*] Agnostic Backoff Triggered... Waiting {delay}s...")
                        time.sleep(delay)
                        # Cap backoff at 32s per architect's directive
                        delay = min(delay * 2, 32)
                        break
                    elif ans == 's':
                        print("[*] Aborting current model. Returning to switch menu...")
                        return None, -1
                    elif ans == 'c':
                        print("[*] Chat context cleared. Aborting current turn...")
                        return None, -2
                    elif ans == 'e':
                        print("[*] Exiting Zero-Shield.")
                        sys.exit(0)
                    else:
                        print("  Invalid input. Type r, s, c, or e.")
            elif "404" in err or "unknown_model" in err.lower():
                next_idx = model_idx + 1
                if next_idx < len(MODEL_REGISTRY):
                    print(f"\n[!] {display} unavailable. Switching to {MODEL_REGISTRY[next_idx][0]}...\n")
                    return call_model(client, next_idx, messages, ui_state)
                print(f"\n[!] {display} unavailable. No fallback models remain. Use /switch.")
                return None, model_idx
            elif "content_filter" in err or "400" in err:
                is_jailbreak = "jailbreak" in err.lower()
                msg = f"\n[!] 400 Bad Request / Content Policy error: {err}"
                if is_jailbreak:
                    msg += ("\n[!] TIP: This was triggered by the provider's 'Jailbreak' filter. "
                            "Avoid aggressive 'ordering' language which triggers safety blocks.")
                print(msg)
                return ("[ORIENT]: Request flagged by content policy or bad request.\n"
                        "[DECIDE]: Returning neutral response.\n"
                        "[ACT]: That request was blocked by the API content policy or format check. "
                        "Please rephrase.", model_idx)
            elif "timeout" in err.lower() or "timed out" in err.lower() or "408" in err.lower():
                print(f"\n[!] {display} suffered an HTTP Timeout (45s). Provider may be overloaded.")
                while True:
                    ans = input(f"  Action: (r)etry immediately, (s)witch model, (c)lear and abort, (e)xit? [r]: ").strip().lower()
                    if not ans or ans == 'r':
                        print("[*] Retrying...")
                        break
                    elif ans == 's':
                        print("[*] Aborting current model. Returning to switch menu...")
                        return None, -1
                    elif ans == 'c':
                        print("[*] Chat context cleared. Aborting current turn...")
                        return None, -2
                    elif ans == 'e':
                        print("[*] Exiting Zero-Shield.")
                        sys.exit(0)
                    else:
                        print("  Invalid input. Type r, s, c, or e.")
            else:
                print(f"\n[!] API Error ({display}): {e}")
                return None, model_idx
                
        attempt += 1
        if attempt >= auto_retries:
            # If auto-retries exhausted, force the user to choose next step
            print(f"\n[!] {display} auto-retries exhausted ({auto_retries}).")
            while True:
                ans = input(f"  Action: (r)etry manually, (s)witch model, (e)xit? [s]: ").strip().lower()
                if ans == 'r':
                    attempt = 0 # reset for manual retry
                    break
                elif not ans or ans == 's':
                    print("[*] Returning to switch menu...")
                    return None, -1
                elif ans == 'e':
                    sys.exit(0)
                else:
                    print("  Invalid input.")


# ─── Main CLI ──────────────────────────────────────────────────────────────────
def run_cli(persist: bool = True):
    """
    Main CLI entry point with comprehensive error handling and resource management.
    
    Features:
    - Multi-model LLM support (gpt-4o-mini, Llama-3.3-70B, Phi-4, DeepSeek-V3, gpt-4o)
    - Persistent Knowledge Graph with AES encryption
    - Real-time AWS API integration with 32 AWS actions
    - Advanced security hardening (credential redaction, input sanitization)
    - OODA loop framework implementation (Observe → Orient → Decide → Act)
    - Comprehensive state management with atomic file operations
    - Thread-safe startup animation and progress indicators
    - Cross-platform compatibility (Windows, Linux, macOS)
    
    Args:
        persist (bool): Enable automatic state persistence on exit/interrupt
        
    Security Features:
    - AKIA/ASIA key redaction with AIDA preservation
    - Path traversal protection
    - Command injection prevention  
    - Unicode attack mitigation
    - Prompt injection blocking
    - System guard loop prevention
    """
    is_tty = sys.stdin.isatty()
    
    # Display enhanced banner
    if is_tty:
        print_banner()
        
        # Quick start guide
        print_section("Quick Start Guide")
        print(f"  {Colors.WHITE}1.{Colors.RESET} Select a model below")
        print(f"  {Colors.WHITE}2.{Colors.RESET} Ask plain-English questions - Zero-Shield calls tools automatically")
        print(f"  {Colors.WHITE}3.{Colors.RESET} Type {Colors.CYAN}/help{Colors.RESET} for all commands\n")
        
        print_section("Special Features (v2.0.0-dev + Security Hardening)")
        print(f"  {Colors.GREEN}✓{Colors.RESET} {Colors.BOLD}Total Recall{Colors.RESET}     - Names auto-indexed into Knowledge Graph")
        print(f"  {Colors.GREEN}✓{Colors.RESET} {Colors.BOLD}IAM Targeting{Colors.RESET}    - Use /target AKIA... for access keys")
        print(f"  {Colors.GREEN}✓{Colors.RESET} {Colors.BOLD}Persistent KG{Colors.RESET}    - Audit findings survive reboots")
        print(f"  {Colors.GREEN}✓{Colors.RESET} {Colors.BOLD}Enhanced HITL{Colors.RESET}    - Full ID re-entry for destructive actions")
        print(f"  {Colors.GREEN}✓{Colors.RESET} {Colors.BOLD}Encrypted State{Colors.RESET}  - Session files protected at rest\n")
        
        print_section("Available Commands")
        commands = [
            ("/switch", "Change LLM model"),
            ("/status", "Show quota and cooldowns"),
            ("/target", "Set active resource"),
            ("/clear", "Clear chat history"),
            ("/export", "Export Knowledge Graph"),
            ("/reset", "Reset all state"),
            ("/help", "Show detailed help"),
            ("exit", "Exit Zero-Shield"),
        ]
        for cmd, desc in commands:
            print(f"  {Colors.CYAN}{cmd:<10}{Colors.RESET} {Colors.DIM}{desc}{Colors.RESET}")
        print()
    else:
        # Non-TTY mode (piped input)
        print("Zero-Shield CLI v2.0.0-dev (security-hardened)")
        print("Copyright (C) 2026 Jeri L3D | JeriSadeuM | MIT License")
        print("Repository: https://github.com/jerisadeumai/zero-shield-cli\n")

    if not run_preflight(): sys.exit(1)

    # Register SIGINT (Ctrl+C) handler — saves state then exits cleanly
    _last_interrupt = 0
    def _sigint(sig, frame):
        nonlocal _last_interrupt
        now = time.time()
        if now - _last_interrupt < 2:
            print_error("Double-interrupt detected. Forcing immediate exit (state may not be saved)")
            os._exit(1)
        _last_interrupt = now
        print_warning("\nInterrupted. Press Ctrl+C again to force-quit")
        if persist:
            state_save()
        else:
            if input("  Save session state? (y/n): ").strip().lower() == 'y':
                state_save()
            else:
                print_info("Session state NOT saved")
        sys.exit(0)
    signal.signal(signal.SIGINT, _sigint)

    model_idx = select_model()
    show_model_tips(model_idx)
    print("")

    # Attempt session restore from previous state file (respecting the manually selected model)
    restored = state_load()
    if restored:
        _session_ctx['model_idx'] = model_idx # Force the new model into the restored context

    client = OpenAI(
        base_url=os.environ.get("GITHUB_MODELS_URL", "https://models.github.ai/inference"),
        api_key=os.environ.get("GITHUB_TOKEN"),
        timeout=45.0
    )

    # Session Knowledge Graph: load from disk (persistent across runs)
    # MUST be loaded BEFORE fetch_snapshot so Total Recall auto-indexing works at boot
    kg = kg_load()

    # RAG: fetch live AWS snapshot at startup — passes kg for Total Recall auto-indexing
    def refresh_knowledge():
        print("[*] Refreshing environment snapshot (Total Recall mode: indexing names)...")
        _snap = fetch_snapshot(kg=kg)
        _map  = build_sg_map(_snap)
        print(f"[*] {_snap}\n")
        return _snap, _map

    # Startup animation for slow boot
    import threading
    
    def startup_animation():
        """Simple rotating animation during startup"""
        chars = "|/-\\"
        i = 0
        while not startup_complete:
            sys.stdout.write(f"\r[*] Initializing Zero-Shield {chars[i % len(chars)]}")
            sys.stdout.flush()
            time.sleep(0.1)
            i += 1
        sys.stdout.write("\r" + " " * 40 + "\r")  # Clear the line
        sys.stdout.flush()
    
    startup_complete = False
    if sys.stdin.isatty():  # Only show animation in interactive mode
        animation_thread = threading.Thread(target=startup_animation, daemon=True)
        animation_thread.start()
    
    snapshot, sg_map = refresh_knowledge()
    startup_complete = True
    
    if sys.stdin.isatty():
        time.sleep(0.2)  # Let animation clear

    # BUG-10 FIX: removed dead id_boot / sg_boot variables (relics of old index-targeting system)
    last_id     = _session_ctx.get('last_id')
    last_sg_id  = _session_ctx.get('last_sg_id')
    last_vpc_id = _session_ctx.get('last_vpc_id')
    last_bucket = _session_ctx.get('last_bucket')
    last_rds_id = _session_ctx.get('last_rds_id')
    last_access_key = _session_ctx.get('last_access_key')

    if not restored:
        print(f"[*] Active target: {last_id}")

    chat_history  = []
    auto_trigger  = False
    action_streak = {}
    format_strikes = 0 # Format Enforcer strike counter
    skip_persist  = False

    while True:
        _session_ctx.update({
            'model_idx':   model_idx,
            'last_id':     last_id,
            'last_sg_id':  last_sg_id,
            'last_vpc_id': last_vpc_id,
            'last_bucket': last_bucket,
            'last_rds_id': last_rds_id,
            'last_access_key': last_access_key,
        })
        
        if not auto_trigger:
            # agent-v2.0.0-alpha: Loop Guard 2.0 - Streak resets ONLY on fresh user input
            action_streak  = {'count': 0, 'total': 0, 'last_action': None, 'history': []}
            format_strikes = 0 # Reset strikes on fresh user input
            
            # Hallucination Guillotine: Inject Active Target alert into the shell prompt
            target_alert = f"\033[31m[ACTIVE TARGET: NONE - EXPLICIT IDENTIFICATION MANDATORY]\033[0m\n" if not last_id else ""
            
            try:
                # Paste Guard 2.0: DISABLED - was triggering false positives on legitimate typing
                # The console log paste detection (below) is sufficient for security
                
                user_input = input(f"\n{target_alert}{ts()} USER ({MODEL_REGISTRY[model_idx][0]})> ").strip()
                # Strip ANSI/VT100 escape sequences (e.g. ^[[A from arrow keys in CloudShell)
                user_input = re.sub(r'(\x1b|\^\[)\[[0-9;]*[A-Za-z]', '', user_input).strip()
            except EOFError:
                print("\n[*] Session terminated.")
                break

            if user_input.lower() in ('exit', 'quit', '/exit'):
                if input("\n  Are you sure you want to exit Zero-Shield? (y/n) [n]: ").strip().lower() != 'y':
                    print("  [*] Exit aborted. Returning to session.")
                    continue
                if persist and not skip_persist:
                    state_save()
                elif not skip_persist:
                    if input("  Save session state? (y/n): ").strip().lower() == 'y':
                        state_save()
                else:
                    print("[*] Session state NOT saved (session was reset).")
                print("[*] Goodbye."); break

            # ── Slash commands ────────────────────────────────────────────
            if user_input.lower() == '/switch':
                model_idx = select_model()
                show_model_tips(model_idx)
                # Ask user whether to carry chat context to the new model
                carry = input("  Carry chat context to new model? (y/n) [n]: ").strip().lower()
                if carry != 'y':
                    chat_history  = []
                    action_streak = {'count': 0, 'total': 0, 'last_action': None, 'history': []}
                    # Reset ALL context IDs so new model starts blind (no KG bleed)
                    last_id     = None
                    last_sg_id  = None
                    last_vpc_id = None
                    last_bucket = None
                    last_rds_id = None
                    last_access_key = None
                    print("  [*] Chat context cleared for new model.")
                else:
                    action_streak = {'count': 0, 'total': 0, 'last_action': None, 'history': []}
                    print("  [*] Chat context carried over to new model.")
                print(""); continue

            if user_input.lower() == '/tips':
                global _show_tips
                _show_tips = not _show_tips
                state = "ENABLED" if _show_tips else "DISABLED"
                print(f"[*] Prompting Tips: {state}")
                continue

            if user_input.lower() == '/status':
                print_quota_table(header="Live LLM Quota Status")
                tok_est = estimate_tokens(chat_history[-14:])
                print(f"\n  Chat context window: ~{tok_est} tokens in last {min(14, len(chat_history))} messages")
                print(f"  Session KG entries : instances={len(kg.get('instances',{}))}, "
                      f"sg_rules={len(kg.get('sg_rules',{}))}, vpcs={len(kg.get('vpcs',{}))}")
                continue

            if user_input.lower() == '/clear':
                if input("\n  Are you sure you want to clear your chat history? (y/n) [n]: ").strip().lower() != 'y':
                    print("  [*] Clear aborted. Returning to session.")
                    continue
                chat_history  = []
                action_streak = {'count': 0, 'total': 0, 'last_action': None, 'history': []}
                print("[*] Chat history cleared. Context reset (KG and snapshot preserved).")
                continue

            if user_input.lower() == '/reset':
                print("\n[!] WARNING: This will permanently delete your session, Knowledge Graph, and context history.")
                if input("  Are you sure you want to perform a hard reset? (y/n) [n]: ").strip().lower() != 'y':
                    print("  [*] Reset aborted. Returning to session.")
                    continue
                # Hard reset: wipe both persistence files + all in-memory context
                # Quota data is kept by default (expensive to re-collect)
                chat_history  = []
                action_streak = {'count': 0, 'total': 0, 'last_action': None, 'history': []}
                last_id     = None
                last_sg_id  = None
                last_vpc_id = None
                last_bucket = None
                last_rds_id = None
                kg.clear()
                kg.update({'instances': {}, 'sg_rules': {}, 'vpcs': {}})
                _cooldown_until.clear()
                _session_ctx.update({k: None for k in _session_ctx})
                
                try:
                    if os.path.exists(STATE_FILE): os.remove(STATE_FILE)
                    if os.path.exists(KG_FILE): os.remove(KG_FILE)
                except Exception: pass
                
                skip_persist = True

                # Ask about quota data separately — it's hard to recollect
                wipe_quota = input("  Also wipe quota data? (y/n) [n]: ").strip().lower()
                if wipe_quota == 'y':
                    _quota_map.clear()
                    print("  [*] Quota data wiped.")
                else:
                    print("  [*] Quota data preserved across reset.")
                # BUG-07 FIX: removed duplicate file-deletion loop
                print("[*] Full reset complete. KG, session state, context IDs wiped.")
                
                # agent-v2.0.0-alpha: Self-Healing Reset — re-run the boot scan so the tool is "warm"
                snapshot, sg_map = refresh_knowledge()
                
                print("[*] Tip: Use /switch to also pick a fresh model if needed.")
                continue

            if user_input.lower() == '/export':
                ts_str   = datetime.now().strftime("%Y%m%d_%H%M%S")
                out_path = os.path.join(os.path.dirname(__file__), f"kg_export_{ts_str}.json")
                try:
                    with open(out_path, 'w', encoding='utf-8') as f:
                        json.dump({'kg': kg, 'quota_map': _quota_map,
                                   'exported_at': ts_str}, f, indent=2)
                    print(f"[*] Session KG + quota map exported to {out_path}")
                except Exception as ex:
                    print(f"[!] Export failed: {ex}")
                continue

            if user_input.lower().startswith('/target'):
                parts = user_input.split()
                if len(parts) < 2:
                    print("  Usage: /target <i-xxxx|sg-xxxx|alias|none>")
                    continue
                val, r_type = resolve_target(parts[1], kg, sg_map)
                if r_type == 'clear':
                    last_id, last_sg_id, last_access_key, last_vpc_id, last_bucket, last_rds_id = (None,) * 6
                    print("[*] Active context cleared.")
                elif r_type == 'instance':
                    last_id = val
                    print(f"[*] Target resolved to instance: {val}")
                elif r_type == 'sg':
                    last_sg_id = val
                    print(f"[*] Target resolved to security group: {val}")
                elif r_type == 'vpc':
                    last_vpc_id = val
                    print(f"[*] Target resolved to VPC: {val}")
                elif r_type == 'subnet':
                    # Subnets shared the VPC context for now
                    print(f"[*] Target resolved to subnet: {val}")
                elif r_type == 'rds':
                    last_rds_id = val
                    print(f"[*] Target resolved to RDS Instance: {val}")
                elif r_type == 'bucket':
                    last_bucket = val
                    print(f"[*] Target resolved to S3 Bucket: {val}")
                elif r_type == 'key':
                    last_access_key = val
                    print(f"[*] Target resolved to IAM Access Key: {val}")
                else:
                    # agent-v2.0.0-alpha: Clear context on failure to prevent stale target leaks
                    last_id, last_sg_id, last_access_key = None, None, None
                    if r_type == 'collision':
                        # val is the matches list now
                        opts = ", ".join(f"{m[0]} ({m[2]})" for m in val)
                        print(f"  [!] Ambiguous input. Options: {opts}")
                    else:
                        print(f"  [!] Could not resolve '{parts[1]}' to an ID or alias. Context cleared.")
                continue

            if user_input.lower() in ('/help', '/?'):
                print("\n  Available slash commands:")
                print("    /switch   — Change LLM model. Prompts to carry or clear chat context.")
                print("    /status   — Show live quota, token usage, and session KG stats")
                print("    /tips     — Toggle the visibility of model-specific prompting tips")
                print("    /target   — Set active resource by exact ID or name alias.")
                print("                 Examples : /target demo  |  /target sg-0abc  |  /target AKIA...  |  /target none")
                print("                 Supports : EC2 instances (i-xxx), security groups (sg-xxx), IAM access keys (AKIA/ASIA...)")
                print("                 Smart    : partial name match runs through the Alias Resolver.")
                print("                            If ambiguous, a collision list is shown — be more specific.")
                print("    /clear    — Clear chat history (keeps model, KG snapshot, and targets)")
                print("    /reset    — Full wipe: KG + session state. Prompts about quota data separately.")
                print("    /export   — Save session KG + quota map to a timestamped JSON file")
                print("    /help     — Show this help message")
                print("    exit      — Terminate the session (saves state based on --persist flag)")
                print()
                print("  * Special features (agent-v2.0.0-alpha):")
                print("    Total Recall  — Instance names are auto-indexed at boot and on every LIST refresh.")
                print("                    '/target demo' or '/target zero' resolves to the real instance ID.")
                print("    KG Persistence— Audit findings (SG rules, VPC data, IAM) persist across sessions.")
                print("    AKIA Targeting— '/target AKIAIOSFODNN7EXAMPLE' loads an IAM key into context.")
                print("                    Then just say 'deactivate it' — the AI triggers the tool for you.")
                print(f"\n  Session mode: {'--persist (auto-save on exit)' if persist else '--no-persist (ask before save)'}")
                continue
                
            if user_input.startswith('/'):
                print(f"  [!] Unknown command: '{user_input}'. Type /help for a list.")
                continue

            if not user_input:
                continue  # never send empty messages — prevents runaway auto-reply loop

            # Paste Guard: Check if input looks like a raw console log paste
            if "USER (" in user_input or "AI:" in user_input or "[OBSERVE]:" in user_input:
                print("\n[!] PASTE DETECTED: You appear to have pasted a previous console log.")
                print("[!] This can cause runaway tool loops and token bloat.")
                
                # --- 1. THE TERMINAL BUFFER FLUSHER (Eat the rest of the paste first) ---
                universal_flush()
                
                # --- 2. TRUE INTERACTIVE CONFIRMATION ---
                ans = input("  Process the first line of this paste? (y/n) [n]: ").strip().lower()
                
                if ans != 'y':
                    print("  [*] Paste discarded. Returning to session.")
                    continue

            skip_persist = False
            action_streak = {}   # Reset streak on new user input
            prefix = f"[Active Target: {last_id} | VPC: {last_vpc_id or 'unknown'}]\n" if last_id else ""
            chat_history.append({"role": "user", "content": f"{prefix}{user_input}"})
        else:
            auto_trigger = False

        api_id   = MODEL_REGISTRY[model_idx][1]
        sys_msg  = build_sys_msg(snapshot, kg, api_id=api_id, last_id=last_id)
        messages = [{"role": "system", "content": sys_msg}] + chat_history[-14:]

        reply, model_idx = call_model(client, model_idx, messages)
        # If user selected (s)witch from the exception handler, hijack the loop to run /switch logic
        if model_idx == -1:
            print("\n")
            new_idx = select_model()
            show_model_tips(new_idx)
            carry = input("  Carry chat context to new model? (y/n) [n]: ").strip().lower()
            
            if carry != 'y':
                # CONTEXT CLEANSE: Preserve the last prompt but wipe the rest
                last_user_msg = None
                if chat_history:
                    raw_msg = chat_history[-1].get("content", "")
                    # STRIP THE ID PREFIX: If it was auto-replayed, remove the [Active Target...] block
                    clean_msg = re.sub(r'^\[Active Target:.*?\]\n', '', raw_msg, flags=re.DOTALL)
                    last_user_msg = {"role": "user", "content": clean_msg}

                chat_history  = []
                if last_user_msg:
                    chat_history.append(last_user_msg)
                action_streak = {}
                # Reset ALL context IDs so new model starts blind
                last_id, last_sg_id, last_vpc_id, last_bucket, last_rds_id, last_access_key = (None,)*6
                _session_ctx.update({
                    'last_id': None, 'last_sg_id': None, 'last_vpc_id': None,
                    'last_bucket': None, 'last_rds_id': None, 'last_access_key': None
                })
                print("  [*] Chat context cleared. Note: Your last prompt (sanitized) will still be sent.")
            else:
                action_streak = {}
                print("  [*] Chat context carried over to new model.")
            
            print("  [*] Re-submitting your prompt to the new model...")
            model_idx = new_idx
            
            # Update the system message for the new model family!
            api_id   = MODEL_REGISTRY[model_idx][1]
            sys_msg  = build_sys_msg(snapshot, kg, api_id=api_id, last_id=last_id)
            messages = [{"role": "system", "content": sys_msg}] + chat_history[-14:]
            
            # RE-FIRE THE API CALL AUTOMATICALLY
            reply, model_idx = call_model(client, model_idx, messages)
            if reply is None:
                continue # If the second attempt fails, loop resets

        # [agent-v2.0.0-alpha] Format Enforcer (Strike System)
        valid_format = all(tag in reply for tag in ["[ORIENT]", "[DECIDE]", "[ACT]"])
        if not valid_format:
            format_strikes += 1
            if format_strikes >= 3:
                print(f"\n[!] FATAL COGNITIVE COLLAPSE: Model '{MODEL_REGISTRY[model_idx][0]}' repeatedly failed OODA formatting.")
                print(f"[!] {format_strikes} consecutive format strikes. Severing execution to prevent infinite API loop.")
                format_strikes = 0
                auto_trigger   = False
                chat_history.append({"role": "assistant", "content": reply})
                continue
            
            print(f"\n[!] FORMAT STRIKE {format_strikes}/3: AI response missing mandatory tags. Forcing clinical retry...")
            # Inject a steering correction into history to force tag adherence
            chat_history.append({"role": "assistant", "content": reply})
            chat_history.append({"role": "user", "content": "[SYSTEM ERROR]: Your response violates the OODA tagging mandate. You MUST use exactly three tags: [ORIENT], [DECIDE], and [ACT]. Respond again with the correct format."})
            auto_trigger = True
            continue
        
        # Success: reset format strikes if format is valid
        if format_strikes > 0:
            print(f"[*] Format recovered on strike {format_strikes}. Proceeding.")
        format_strikes = 0

        # ESCAPE HATCH (c): Clear context and abort to main prompt
        if model_idx == -2:
            chat_history  = []
            action_streak = {}
            format_strikes = 0
            last_id, last_sg_id, last_vpc_id, last_bucket, last_rds_id, last_access_key = (None,)*6
            _session_ctx.update({
                'last_id': None, 'last_sg_id': None, 'last_vpc_id': None,
                'last_bucket': None, 'last_rds_id': None, 'last_access_key': None
            })
            print("  [*] Chat context cleared. Returning to main prompt.\n")
            auto_trigger = False
            continue

        if reply is None: continue

        # [agent-v2] Double-Layer Redaction: Scrub AI output before it hits console or history
        reply = _redact_secrets(reply)

        print(f"{ts()} AI: {reply}\n")
        chat_history.append({"role": "assistant", "content": reply})

        # [agent-v2] Explicit Targeting: Removed automatic ID scraping from conversational context.
        # last_id, last_sg_id, etc. are now ONLY updated via /target or [ACTION:TARGET].

        # Detect action
        action, action_param = detect_action(reply)
        observation = None

        # --- Action loop guard 2.0 ---
        if action:
            # agent-v2.0.0-alpha: Lazy Shield - ensure streak is initialized if it was wiped by a reset
            if not isinstance(action_streak, dict) or 'history' not in action_streak:
                action_streak = {'count': 0, 'total': 0, 'last_action': None, 'history': []}
            if action == action_streak.get('last_action'):
                action_streak['count'] = action_streak.get('count', 0) + 1
            else:
                action_streak['last_action'] = action
                action_streak['count'] = 1

            # Circular Loop Detection: Catch non-consecutive repeating actions
            history_key = f"{action}:{action_param or ''}"
            action_streak['history'].append(history_key)
            if action_streak['history'].count(history_key) >= 3:
                observation = (f"[SYSTEM GUARD]: Action '{history_key}' attempted 3 times in this turn. "
                               f"Potential circular loop. Forcing manual override.")
                print(f"\n[!] {observation}")
                action = None
                auto_trigger = False
                continue

            if action_streak['count'] >= 3:
                observation = (f"[SYSTEM GUARD]: [{action}] called {action_streak['count']} times consecutively. "
                               f"Cognitive loop detected. Forcing manual override.")
                print(f"\n[!] {observation}")
                action_streak = {'count': 0, 'total': 0, 'last_action': None, 'history': []} # Reset
                action = None
                auto_trigger = False
                continue

            action_streak['total'] = action_streak.get('total', 0) + 1
            if action_streak['total'] > 12:
                observation = ("[SYSTEM GUARD]: Maximum automated tool chain exceeded (12 turns). "
                               "Please summarize findings and return control to the user.")
                print(f"\n[!] {observation}")
                action_streak['total'] = 0
                action = None
                auto_trigger = False
        else:
            # agent-v2.0.0-alpha: Don't wipe the total streak just because AI was chatty.
            # Only reset the consecutive count.
            action_streak['last_action'] = None
            action_streak['count'] = 0

        # --- Dispatch tool ---
        if action == "MULTIPLE_ACTIONS_DETECTED":
            observation = "[SYSTEM GUARD \u26a0] You attempted to trigger multiple actions in one turn. This violates OODA loop physics. Execute strictly ONE [ACTION] tag, wait for [OBSERVE], then proceed."
            print(f"\n[!] {observation}")

        elif action == "LIST":
            observation = tool_list_resources(kg=kg)
            kg_save(kg)
            snapshot = f"Instances: {observation}\n" + "\n".join(snapshot.split("\n")[1:])
            print(f"[OBSERVE]: {observation}")

        elif action == "LIST_SGS":
            observation = tool_list_security_groups()
            sg_map = build_sg_map(observation)
            kg_save(kg)
            snapshot = snapshot.split("\n")[0] + f"\n{observation}"
            print(f"[OBSERVE]: {observation}")

        elif action == "INSPECT":
            target = action_param or last_id
            val, r_type = resolve_target(target, kg, sg_map)
            if r_type == 'instance' or (val and isinstance(val, str) and val.startswith('i-')):
                observation = tool_inspect_resource(val)
                last_id     = val
                kg['instances'][val] = observation
                kg_save(kg)
                # agent-v2.0.0-alpha: Context Unification - Auto-warm the VPC context from Instance metadata
                _vpc_match = re.search(r'VPC:\s*(vpc-[a-z0-9]+)', observation)
                if _vpc_match: last_vpc_id = _vpc_match.group(1)
                print(f"[OBSERVE]: {observation}")
            elif r_type == 'collision':
                opts = "\n".join(f"- {m[0]} ({m[2]})" for m in val)
                observation = (f"ERROR: Ambiguous target '{target}'. Matches found:\n{opts}\n"
                               "Please ask the user to be more specific or choose one of these IDs.")
                print(f"[!] {observation}")
            else:
                # val is the list of all known aliases for suggestions
                suggestions = ", ".join(str(v) for v in val[:5]) if isinstance(val, list) else ""
                observation = f"Error: '{target}' not found. Did you mean one of these: {suggestions}?"
                print(f"[!] {observation}")

        elif action == "SG_RULES":
            target = action_param or last_sg_id
            val, r_type = resolve_target(target, kg, sg_map) if target else (None, None)
            if r_type == 'sg' or (val and val.startswith('sg-')):
                observation = tool_sg_rules(val)
                last_sg_id  = val
                kg['sg_rules'][val] = observation
                kg_save(kg)
                print(f"[OBSERVE]: {observation}")
                hints = auto_remediation_hint(observation)
                if hints: print(f"\n{hints}")
                ground_truth = interpret_sg_rules(observation)
                observation = observation + "\n\n" + ground_truth
            else:
                observation = f"Error: '{target}' is not a valid security group ID or alias."
                print(f"[!] {observation}")

        elif action == "VPC_INFO":
            target = action_param or last_vpc_id
            
            # agent-v2.0.0-alpha: Explicit Resolve - if a name like 'demo' is used, resolve it to a VPC
            _val, _type = resolve_target(target, kg, sg_map) if target else (None, None)
            if _type == 'instance' or (_val and str(_val).startswith('i-')):
                # Adopting from a named instance
                _inst_data = kg.get('instances', {}).get(_val, '')
                _v_hit = re.search(r'VPC:\s*(vpc-[a-z0-9]+)', _inst_data)
                if _v_hit: target = _v_hit.group(1)
            elif _val and str(_val).startswith('vpc-'):
                target = _val

            # Fallback: Implicit adoption from active instance context
            if not target and last_id:
                _inst_data = kg.get('instances', {}).get(last_id, '')
                _v_hit = re.search(r'VPC:\s*(vpc-[a-z0-9]+)', _inst_data)
                if _v_hit: target = _v_hit.group(1)

            if target:
                observation = tool_vpc_info(target)
                last_vpc_id = target
                kg['vpcs'][target] = observation
                kg_save(kg)
                print(f"[OBSERVE]: {observation}")
            else:
                observation = ("Error: No VPC ID specified. Name a target (e.g. 'demo') or use INSPECT.")
                print(f"[!] {observation}")

        # BUG-01 FIX: removed zombie TARGET block that was here (lines 1955-1961 in old file).
        # It hard-assigned last_id = action_param with ZERO resolver logic, and preempted
        # the real resolver-backed TARGET block below, making it permanently unreachable.

        elif action == "LOGS":
            target = action_param or last_id
            val, r_type = resolve_target(target, kg, sg_map) if target else (None, None)
            if r_type == 'instance' or (val and val.startswith('i-')):
                observation = tool_get_logs(val)
                last_id     = val
                print(f"[OBSERVE]: {observation}")
            else:
                observation = f"Error: '{target}' is not a valid instance ID or alias."
                print(f"[!] {observation}")

        elif action == "CW_LOGS":
            if last_id:
                observation = tool_cw_logs(last_id)
                print(f"[OBSERVE]: {observation}")
            else:
                observation = "Error: No instance ID in context. Use LIST first."
                print(f"[!] {observation}")

        elif action == "IAM_CHECK":
            if last_id:
                observation = tool_iam_check(last_id)
                kg['instances'][last_id] = kg.get('instances', {}).get(last_id, '') + f"\n[IAM] {observation}"
                kg_save(kg)
                print(f"[OBSERVE]: {observation}")
            else:
                observation = "Error: No instance ID in context. Use LIST first."
                print(f"[!] {observation}")

        elif action == "COST":
            if last_id:
                observation = tool_cost_insight(last_id)
                print(f"[OBSERVE]: {observation}")
            else:
                observation = "Error: No instance ID in context. Use LIST first."
                print(f"[!] {observation}")

        elif action == "MODIFY_SG":
            # PYTHON HARD GATE: Block execution if no active target is set
            if not last_id:
                observation = "[PYTHON GATE TRIGGERED] MODIFY_SG blocked: No active target. Use /target <id> to set a target first."
                print(f"\n[!] {observation}")
                chat_history.append({"role": "user", "content": f"[OBSERVE]: {observation}"})
                auto_trigger = True
            elif last_id and last_sg_id:
                sg_display = next((n for n, i in sg_map.items() if i == last_sg_id), last_sg_id)
                # CRITICAL-04 FIX: Enhanced HITL with instance ID re-entry
                print(f"\n[HITL] CRITICAL ACTION: Move {last_id} → {last_sg_id} ({sg_display})?")
                print(f"  This will change the security group for instance {last_id}.")
                print(f"  To confirm, type the instance ID: {last_id}")
                
                confirmation = input("  Enter instance ID to confirm: ").strip().lower()
                
                if confirmation == last_id.lower():
                    time.sleep(1)  # Brief delay to prevent accidental rapid confirmations
                    observation = tool_modify_sg(last_id, last_sg_id)
                    snapshot = fetch_snapshot(kg=kg)  # BUG-04 FIX: pass kg so TotalRecall stays live
                    sg_map   = build_sg_map(snapshot)
                    print(f"[ACT]: {observation}")
                else:
                    observation = f"User cancelled SG modification (confirmation mismatch: '{confirmation}' != '{last_id}')."
                    print(f"[OBSERVE]: {observation}")
            else:
                missing      = "instance ID" if not last_id else "target SG ID"
                observation  = f"Error: Cannot MODIFY_SG — {missing} not in context."
                print(f"[!] {observation}")

        elif action == "QUARANTINE":
            # PYTHON HARD GATE: Block execution if no active target is set
            if not last_id:
                observation = "[PYTHON GATE TRIGGERED] QUARANTINE blocked: No active target. Use /target <id> to set a target first."
                print(f"\n[!] {observation}")
                chat_history.append({"role": "user", "content": f"[OBSERVE]: {observation}"})
                auto_trigger = True
            else:
                # CRITICAL-04 FIX: Enhanced HITL with instance ID re-entry
                print(f"\n[HITL] CRITICAL ACTION: Quarantine {last_id} with SG {QUARANTINE_SG_ID}?")
                print(f"  This will ISOLATE the instance by changing its security group.")
                print(f"  To confirm, type the instance ID: {last_id}")
                
                confirmation = input("  Enter instance ID to confirm: ").strip().lower()
                
                if confirmation == last_id.lower():
                    time.sleep(1)  # Brief delay to prevent accidental rapid confirmations
                    observation = tool_quarantine(last_id)
                    snapshot = fetch_snapshot(kg=kg)  # BUG-04 FIX: pass kg so TotalRecall stays live
                    sg_map   = build_sg_map(snapshot)
                    print(f"[ACT]: {observation}")
                else:
                    observation = f"User aborted quarantine (confirmation mismatch: '{confirmation}' != '{last_id}')."
                    print(f"[OBSERVE]: {observation}")

        elif action == "EC2_VOLUMES":
            observation = tool_ec2_volumes()
            print(f"[OBSERVE]: {observation}")

        elif action == "EC2_SNAPSHOTS":
            observation = tool_ec2_snapshots()
            print(f"[OBSERVE]: {observation}")

        elif action == "EC2_KEYPAIRS":
            observation = tool_ec2_keypairs()
            print(f"[OBSERVE]: {observation}")

        elif action == "EC2_NACLS":
            target = action_param or last_vpc_id
            
            # agent-v2.0.0-alpha: Explicit Resolve support for aliases
            _val, _type = resolve_target(target, kg, sg_map) if target else (None, None)
            if _type == 'instance' or (_val and str(_val).startswith('i-')):
                _inst_data = kg.get('instances', {}).get(_val, '')
                _v_hit = re.search(r'VPC:\s*(vpc-[a-z0-9]+)', _inst_data)
                if _v_hit: target = _v_hit.group(1)
            elif _val and str(_val).startswith('vpc-'):
                target = _val

            # Fallback: Adoption from Active Instance Context
            if not target and last_id:
                _inst_data = kg.get('instances', {}).get(last_id, '')
                _v_hit = re.search(r'VPC:\s*(vpc-[a-z0-9]+)', _inst_data)
                if _v_hit: target = _v_hit.group(1)

            if target:
                observation = tool_ec2_nacls(target)
                last_vpc_id = target
                print(f"[OBSERVE]: {observation}")
            else:
                observation = "Error: No VPC ID specified. Target an instance or VPC alias first."
                print(f"[!] {observation}")

        elif action == "IAM_USERS":
            observation = tool_iam_users()
            kg['iam_users'] = observation
            kg_save(kg)
            print(f"[OBSERVE]: {observation}")

        elif action == "IAM_ROLES":
            observation = tool_iam_roles()
            print(f"[OBSERVE]: {observation}")

        elif action == "IAM_KEYS":
            observation = tool_iam_keys()
            print(f"[OBSERVE]: {observation}")

        elif action == "DEACTIVATE_ACCESS_KEY":
            target = action_param or last_access_key
            val, r_type = resolve_target(target, kg, sg_map) if target else (None, None)
            if r_type == 'key' or (val and (val.startswith('AKIA') or val.startswith('ASIA')) and len(val)==20):
                # CRITICAL-04 FIX: Enhanced HITL with resource ID re-entry and rate limiting
                print(f"\n[HITL] CRITICAL ACTION: Deactivate Identity Access Key {val}?")
                print(f"  This will IMMEDIATELY revoke access for this credential.")
                print(f"  To confirm, type the FULL access key ID: {val}")
                
                confirmation = input("  Enter key ID to confirm: ").strip().upper()
                
                if confirmation == val:
                    # Add a brief delay to prevent accidental rapid confirmations
                    time.sleep(1)
                    observation = tool_deactivate_access_key(val)
                    last_access_key = val
                    print(f"[ACT]: {observation}")
                else:
                    observation = f"User aborted key deactivation (confirmation mismatch: '{confirmation}' != '{val}')."
                    print(f"[OBSERVE]: {observation}")
            else:
                observation = f"Error: '{target}' is not a valid IAM access key. Use IAM_KEYS first."
                print(f"[!] {observation}")

        elif action == "S3_LIST":
            observation = tool_s3_list()
            kg['s3_buckets'] = observation
            kg_save(kg)
            # Extract first real bucket name (skip header line 'S3 Buckets:')
            for _line in observation.splitlines():
                _bkt = _line.strip()
                if _bkt and not _bkt.startswith('S3') and '|' in _bkt:
                    last_bucket = _bkt.split('|')[0].strip()
                    break
            print(f"[OBSERVE]: {observation}")

        elif action == "S3_POLICY":
            target = action_param or last_bucket
            if target:
                observation = tool_s3_policy(target)
                last_bucket = target
                print(f"[OBSERVE]: {observation}")
            else:
                observation = "Error: No bucket specified. Use S3_LIST first."
                print(f"[!] {observation}")

        elif action == "RDS_LIST":
            observation = tool_rds_list()
            kg['rds'] = observation
            kg_save(kg)
            # Extract first real RDS identifier (skip header line 'RDS Instances:')
            for _line in observation.splitlines():
                _rds = _line.strip()
                if _rds and not _rds.startswith('RDS') and '|' in _rds:
                    last_rds_id = _rds.split('|')[0].strip()
                    break
            print(f"[OBSERVE]: {observation}")

        elif action == "LAMBDA_LIST":
            observation = tool_lambda_list()
            kg['lambda'] = observation
            kg_save(kg)
            print(f"[OBSERVE]: {observation}")

        elif action == "CW_ALARMS":
            observation = tool_cw_alarms()
            print(f"[OBSERVE]: {observation}")

        elif action == "CW_METRICS":
            target = action_param or last_id
            val, r_type = resolve_target(target, kg, sg_map) if target else (None, None)
            if r_type == 'instance' or (val and val.startswith('i-')):
                observation = tool_cw_metrics(val)
                last_id     = val
                print(f"[OBSERVE]: {observation}")
            else:
                observation = f"Error: '{target}' is not a valid instance."
                print(f"[!] {observation}")

        elif action == "CLOUDTRAIL":
            observation = tool_cloudtrail()
            kg['trail'] = observation
            kg_save(kg)
            print(f"[OBSERVE]: {observation}")

        elif action == "COST_EXPLORER":
            observation = tool_cost_explorer()
            kg['cost'] = observation
            kg_save(kg)
            print(f"[OBSERVE]: {observation}")

        elif action == "KMS_KEYS":
            observation = tool_kms_keys()
            kg['kms'] = observation
            kg_save(kg)
            print(f"[OBSERVE]: {observation}")

        elif action == "DYNAMODB_LIST":
            observation = tool_dynamodb_list()
            kg['dynamo'] = observation
            kg_save(kg)
            print(f"[OBSERVE]: {observation}")

        elif action == "EFS_LIST":
            observation = tool_efs_list()
            kg['efs'] = observation
            kg_save(kg)
            print(f"[OBSERVE]: {observation}")

        elif action == "WAF_WEBACLS":
            observation = tool_waf_webacls()
            kg['waf'] = observation
            kg_save(kg)
            print(f"[OBSERVE]: {observation}")

        elif action == "TARGET":
            val, r_type = resolve_target(action_param, kg, sg_map)
            
            # F1 Patch: Hard Crosshair Reset. Wipe ALL state before assigning a new singular focus.
            last_id, last_sg_id, last_access_key, last_vpc_id, last_bucket, last_rds_id = (None,) * 6
            
            if r_type == 'clear':
                observation = "SUCCESS: Context target cleared."
            elif r_type == 'instance':
                last_id = val
                observation = f"SUCCESS: Active target set to instance {val}."
            elif r_type == 'sg':
                last_sg_id = val
                observation = f"SUCCESS: Active target set to security group {val}."
            elif r_type == 'vpc':
                last_vpc_id = val
                observation = f"SUCCESS: Active target set to VPC {val}."
            elif r_type == 'rds':
                last_rds_id = val
                observation = f"SUCCESS: Active target set to RDS instance {val}."
            elif r_type == 'bucket':
                last_bucket = val
                observation = f"SUCCESS: Active target set to S3 Bucket {val}."
            elif r_type == 'key':
                last_access_key = val
                observation = f"SUCCESS: Active target set to access key {val}."
            else:
                if r_type == 'collision':
                    opts = "\n".join(f"- {m[0]} ({m[2]})" for m in val)
                    observation = (f"ERROR: Ambiguous target '{action_param}'. Matches:\n{opts}\n"
                                   "Ask the user for clarification.")
                else:
                    suggestions = ", ".join(str(v) for v in val[:5]) if isinstance(val, list) else ""
                    observation = f"ERROR: Could not resolve '{action_param}'. Known aliases: {suggestions}"
            print(f"[*] {observation}")

        elif action == "GUARDDUTY":
            observation = tool_guardduty()
            print(f"[OBSERVE]: {observation}")

        if observation:
            safe_obs = _redact_secrets(observation)
            family = MODEL_FAMILIES.get(MODEL_REGISTRY[model_idx][1], "gpt")
            
            if family == "gpt":
                # Safe fallback to prevent Azure Jailbreak Filter flags
                content_str = f"Tool output received:\n{safe_obs}"
            else:
                # Strong reinforcement for Llama/Phi/DeepSeek to prevent looping
                content_str = (f"[SYSTEM OBSERVATION FEEDBACK]\n{safe_obs}\n\n"
                               "Summarize this result to directly answer the user's request. "
                               "Do not trigger another tool unless genuinely more data is needed.")
                               
            chat_history.append({
                "role": "user",
                "content": content_str
            })
            auto_trigger = True

# ─── Per-model Prompting Tips ─────────────────────────────────────────────────
MODEL_TIPS = {
    "gpt": [
        "Works best with direct, specific requests.",
        "Handles multi-step chains well (e.g. 'Inspect, then check SG rules').",
        "Tip: Avoid raw key/credential content in chat — may trigger content filters.",
        "Supports structured audit requests: 'Give me a full security posture summary'.",
    ],
    "llama": [
        "Keep requests focused and single-topic per turn.",
        "Warning: 8B models may guess; demand 'verified metadata'.",
        "E.g.: 'Audit VPC isolation for [1]. First call INSPECT,",
        "then VPC_INFO for CIDRs. I need verified data.'",
    ],
    "phi": [
        "Very instruction-compliant — phrase requests as clear commands.",
        "If it outputs two [ACTION:] tags at once, that is a known quirk. Press Enter to continue.",
        "Tip: Ask one thing at a time for most reliable results.",
        "Best for: SG rules audits and targeted single-resource queries.",
    ],
    "deepseek": [
        "Excellent at multi-step correlative reasoning (logs + SG rules cross-analysis).",
        "Be concise — it tends to be verbose, so ask for 'a brief summary' if needed.",
        "Tip: Subject to stricter rate limits. If 429 appears, switch to gpt-4o-mini.",
        "Best for: threat analysis, network topology, and detective-style queries.",
    ],
}

def show_model_tips(model_idx):
    if not _show_tips: return
    api_id = MODEL_REGISTRY[model_idx][1]
    family = MODEL_FAMILIES.get(api_id, "gpt")
    tips   = MODEL_TIPS.get(family, MODEL_TIPS["gpt"])
    display = MODEL_REGISTRY[model_idx][0]
    print(f"\n┌─ Prompting Tips for {display} {'─' * max(0, 36-len(display))}┐")
    for tip in tips:
        # Wrap tip to 53 chars
        print(f"│  • {tip:<53}│")
    print("└─────────────────────────────────────────────────────────┘")


def print_quota_table(header: str = "Select LLM for Inference"):
    """Print the model selection table with live quota data when available."""
    from datetime import datetime as _dt
    now = _dt.now().timestamp()
    w = 84
    print(f"\n┌─ {header} " + "─" * max(0, w - 3 - len(header)) + "┐")
    for idx, (display, _, _, _, desc) in enumerate(MODEL_REGISTRY):
        q = _quota_map.get(idx, {})
        reset_ts = _cooldown_until.get(idx)
        wait_s = max(0, int((reset_ts or now) - now))
        rem_req = q.get('remaining_req')
        
        if wait_s > 0:
            m2, s2 = divmod(wait_s, 60)
            state = f"COOLDOWN (~{m2}m {s2}s)"
        elif rem_req is not None and rem_req == 0:
            # Cooldown passed, but last known quota was 0.
            state = "REPLENISHING?"
        else:
            state = "AVAILABLE"
            
        line1 = f"[{idx+1}] {display:<28} | {desc}"
        print(f"│ {line1:<{w-2}} │")
        
        if not q.get('verified'):
            line2 = f"      State  : {state}"
            line3 = f"      (quota unverified - awaiting first API request to fetch HTTP headers)"
            print(f"│ {line2:<{w-2}} │")
            print(f"│ {line3:<{w-2}} │")
        else:
            limit_type = q.get('limit_type', '')
            lt_str = f"| Trigger: {limit_type}" if (state != "AVAILABLE" and limit_type) else ""
            
            # Show exact tokens used
            tokens = q.get('total_tokens', 0)
            token_str = f"| Tokens Used: {tokens:<6}" if tokens > 0 else ""
            
            line2 = f"      State  : {state:<25} {lt_str} {token_str}"
            print(f"│ {line2:<{w-2}} │")
            
            rl = _reset_label(idx)
            if rl:
                line_rl = f"      Reset  : {rl}"
                print(f"│ {line_rl:<{w-2}} │")
                
            req_bar = _quota_req_bar(idx, 20)
            if req_bar:
                line_req = f"      Reqs   : {req_bar}"
                print(f"│ {line_req:<{w-2}} │")
                
            tok_bar = _quota_tok_bar(idx, 20)
            if tok_bar:
                line_tok = f"      Toks   : {tok_bar}"
                print(f"│ {line_tok:<{w-2}} │")
                
        if idx < len(MODEL_REGISTRY) - 1:
            print("├" + "─" * w + "┤")
    print("└" + "─" * w + "┘")
    print("  * Note: If a model hides its HTTP Quota Headers, exact limits will display as 'Unknown'.\n")

def select_model():
    print_quota_table()
    from datetime import datetime as _dt
    while True:
        choice = input(f"  Model (1-{len(MODEL_REGISTRY)}): ").strip()
        if choice.lower() in ('exit', 'quit', '/exit'):
            print("\n[*] Exiting Zero-Shield.")
            import sys
            sys.exit(0)
            
        if choice.isdigit() and 1 <= int(choice) <= len(MODEL_REGISTRY):
            idx = int(choice) - 1
            
            # THE SKEPTICAL CHOICE IMPLEMENTATION
            now = _dt.now().timestamp()
            reset_ts = _cooldown_until.get(idx, 0)
            wait_s = max(0, int(reset_ts - now))
            
            if wait_s > 0:
                print(f"\n[!] WARNING: {MODEL_REGISTRY[idx][0]} is currently on a {wait_s}s cooldown.")
                print("    Selecting it now will likely result in an immediate API failure.")
                if input("    Are you sure you want to force this selection? (y/n) [n]: ").strip().lower() != 'y':
                    print("    [*] Selection aborted. Please choose an available model.")
                    continue
            
            return idx
            
        print("  Invalid selection.")

if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser(
        description="Zero-Shield: Agentic AWS Security Copilot",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog=(
            "Session persistence flags:\n"
            "  --persist      Auto-save context on exit/Ctrl+C and prompt to restore on startup (DEFAULT).\n"
            "  --no-persist   Never auto-save. You will be asked before each save."
        )
    )
    persist_group = parser.add_mutually_exclusive_group()
    persist_group.add_argument('--persist',    dest='persist', action='store_true',  default=True,
                               help='Auto-save session on exit (default).')
    persist_group.add_argument('--no-persist', dest='persist', action='store_false',
                               help='Ask before saving; never auto-save.')
    parser.add_argument('--no-tips', dest='show_tips', action='store_false', default=True,
                        help='Hide model-specific prompting tips at startup.')
    args = parser.parse_args()
    
    _show_tips = args.show_tips
    
    run_cli(persist=args.persist)