cloud_soc_env.py

"""
OpenEnv-CloudSOC: Cloud Security Operations Center Benchmark Environment
=========================================================================
A Gymnasium-compatible environment for evaluating LLM agents on cloud security
incident response. Implements Zero-DB architecture with in-memory state management.

Hardware Target: 2 vCPU / 8 GB RAM Docker container
Architecture: Thin Client, Deep State (no external databases)

Key Features:
- Realistic AWS cloud simulation with deceptive elements
- Progressive reward shaping through forensic kill-chain
- Strict causal dependencies for proper IR procedure
- Adversarial traps to test agent judgment
- Multi-task campaign with state persistence
- Comprehensive audit trail and analytics
"""

import random
import json
import hashlib
import copy
from datetime import datetime, timedelta
from typing import Dict, List, Optional, Any, Tuple, Union, Set
from enum import Enum
from dataclasses import dataclass, field
from pydantic import BaseModel, Field, validator
from collections import defaultdict
import gymnasium as gym
from gymnasium import spaces

class Observation(BaseModel):
    state_description: str
    metadata: Optional[Dict[str, Any]] = None

class Action(BaseModel):
    tool: str
    args: Dict[str, Any]

class Reward(BaseModel):
    value: float
    reason: Optional[str] = None


# =============================================================================
# CONSTANTS AND CONFIGURATION
# =============================================================================

# Reward constants for fine-tuned scoring
REWARDS = {
    "FLAG_DISCOVERY": 0.02,
    "CRITICAL_ACTION": 0.10,
    "CONTAINMENT_ACTION": 0.05,
    "PARSE_ERROR": -0.02,
    "PRECONDITION_FAIL": -0.05,
    "EVIDENCE_DESTRUCTION": -1.0,
    "QUERY_BASIC": -0.01,
    "QUERY_DEEP": -0.05,
    "INVALID_TOOL": -0.01,
    "TIMELINE_BONUS_MAX": 0.5,
}

# Malicious IPs for realistic simulation
MALICIOUS_IPS = [
    "185.220.101.42",  # TOR exit node
    "45.155.205.233",  # Known C2
    "194.26.192.64",   # Bulletproof hosting
]

# Benign IPs for noise
BENIGN_IPS = [
    "10.0.0.1", "10.0.1.50", "10.0.5.20", "172.16.0.100",
    "192.168.1.10", "10.0.0.254", "172.31.0.1"
]


# =============================================================================
# PYDANTIC MODELS FOR TOOL ABSTRACTION (Mechanic #7)
# =============================================================================

class ToolArgs(BaseModel):
    """Base class for tool arguments"""
    class Config:
        extra = "forbid"


class CloudWatchQueryArgs(ToolArgs):
    log_group: str = Field(..., description="CloudWatch log group to query")
    start_time: Optional[str] = Field(None, description="ISO timestamp for query start")
    end_time: Optional[str] = Field(None, description="ISO timestamp for query end")
    filter_pattern: Optional[str] = Field(None, description="CloudWatch filter pattern")


class EC2ActionArgs(ToolArgs):
    instance_id: str = Field(..., pattern=r"^i-[a-f0-9]+$")


class IAMRoleArgs(ToolArgs):
    role_name: str = Field(..., min_length=1, max_length=64)


class IAMDetachArgs(ToolArgs):
    role_name: str
    instance_id: Optional[str] = None


class S3BucketArgs(ToolArgs):
    bucket_name: str = Field(..., pattern=r"^[a-z0-9][a-z0-9.-]{1,61}[a-z0-9]$")


class CredentialRotateArgs(ToolArgs):
    resource_type: str = Field(..., description="Type: 'rds', 'iam', 's3'")
    resource_id: str


class SnapshotArgs(ToolArgs):
    instance_id: str = Field(..., pattern=r"^i-[a-f0-9]+$")
    snapshot_type: str = Field(default="forensic", description="'forensic' or 'backup'")


class IncidentCloseArgs(ToolArgs):
    incident_id: str
    timeline: List[str] = Field(..., min_items=1, description="Ordered list of incident events")
    root_cause: Optional[str] = None
    remediation_summary: Optional[str] = None


class SecurityGroupArgs(ToolArgs):
    security_group_id: str = Field(..., pattern=r"^sg-[a-f0-9]+$")
    action: str = Field(..., description="'isolate' or 'restore'")


class InvestigateArgs(ToolArgs):
    target_type: str = Field(..., description="'ip', 'user', 'role', 'instance'")
    target_id: str


class ToolCall(BaseModel):
    """Strict tool call schema that agents must use"""
    thought: str = Field(..., description="1-sentence internal reasoning", max_length=500)
    tool: str = Field(..., description="Tool name to invoke")
    args: Dict[str, Any] = Field(default_factory=dict)

    @validator('tool')
    def validate_tool_name(cls, v):
        valid_tools = {
            'aws.cloudwatch.query_basic', 'aws.cloudwatch.query_deep',
            'aws.ec2.isolate', 'aws.ec2.terminate', 'aws.ec2.describe',
            'aws.ec2.snapshot', 'aws.iam.detach_role', 'aws.iam.revoke_credentials',
            'aws.iam.describe_role', 'aws.iam.list_policies',
            'aws.s3.get_bucket_policy', 'aws.s3.block_public_access',
            'aws.s3.list_objects', 'aws.rds.rotate_credentials',
            'aws.security_group.modify', 'aws.investigate',
            'aws.soc.close_incident', 'aws.soc.get_alerts',
            'aws.guardduty.get_findings', 'aws.cloudtrail.lookup_events',
            'aws.config.get_compliance', 'aws.ssm.run_command',
            'aws.lambda.list_functions', 'aws.sts.get_caller_identity'
        }
        if v not in valid_tools:
            raise ValueError(f"Unknown tool: {v}. Valid tools: {sorted(valid_tools)}")
        return v
    
    @validator('thought')
    def validate_thought(cls, v):
        if not v or len(v.strip()) < 5:
            raise ValueError("Thought must contain meaningful reasoning (min 5 chars)")
        return v.strip()


# =============================================================================
# STATE MODELS (In-Memory Cloud State)
# =============================================================================

class InstanceState(Enum):
    RUNNING = "running"
    STOPPED = "stopped"
    ISOLATED = "isolated"
    TERMINATED = "terminated"


class IncidentPhase(Enum):
    INVESTIGATION = "investigation"
    CONTAINMENT = "containment"
    ERADICATION = "eradication"
    RECOVERY = "recovery"
    CLOSED = "closed"


@dataclass
class EC2Instance:
    instance_id: str
    state: InstanceState = InstanceState.RUNNING
    attached_roles: List[str] = field(default_factory=list)
    security_groups: List[str] = field(default_factory=list)
    is_compromised: bool = False
    has_forensic_snapshot: bool = False
    metadata: Dict[str, Any] = field(default_factory=dict)


@dataclass
class IAMRole:
    role_name: str
    policies: List[str] = field(default_factory=list)
    is_compromised: bool = False
    is_detached: bool = False
    credentials_rotated: bool = False
    has_backdoor: bool = False


@dataclass
class S3Bucket:
    bucket_name: str
    is_public: bool = False
    contains_credentials: bool = False
    leaked_keys: List[str] = field(default_factory=list)
    public_access_blocked: bool = False


@dataclass
class LogEntry:
    timestamp: str
    source: str
    message: str
    severity: str  # INFO, WARN, ERROR, CRITICAL
    is_attack_indicator: bool = False
    is_noise: bool = False
    is_red_herring: bool = False


@dataclass 
class Alert:
    alert_id: str
    title: str
    severity: str
    source: str
    timestamp: str
    details: Dict[str, Any] = field(default_factory=dict)
    is_true_positive: bool = True


@dataclass
class CloudState:
    """Complete in-memory cloud infrastructure state"""
    instances: Dict[str, EC2Instance] = field(default_factory=dict)
    roles: Dict[str, IAMRole] = field(default_factory=dict)
    buckets: Dict[str, S3Bucket] = field(default_factory=dict)
    security_groups: Dict[str, Dict] = field(default_factory=dict)
    logs: List[LogEntry] = field(default_factory=list)
    alerts: List[Alert] = field(default_factory=list)
    credentials: Dict[str, Dict] = field(default_factory=dict)
    
    # Investigation state
    discovered_flags: set = field(default_factory=set)
    completed_actions: List[str] = field(default_factory=list)
    incident_phase: IncidentPhase = IncidentPhase.INVESTIGATION
    
    # Timeline tracking for grading
    agent_timeline: List[str] = field(default_factory=list)
    ground_truth_timeline: List[str] = field(default_factory=list)
    
    # Enhanced tracking
    evidence_collected: Dict[str, Any] = field(default_factory=dict)
    mitre_techniques: Set[str] = field(default_factory=set)
    ioc_indicators: List[Dict] = field(default_factory=list)
    
    def to_dict(self) -> Dict:
        """Serialize state for debugging/checkpointing"""
        return {
            "instances": {k: self._inst_to_dict(v) for k, v in self.instances.items()},
            "roles": {k: self._role_to_dict(v) for k, v in self.roles.items()},
            "buckets": {k: self._bucket_to_dict(v) for k, v in self.buckets.items()},
            "discovered_flags": list(self.discovered_flags),
            "completed_actions": self.completed_actions,
            "phase": self.incident_phase.value,
            "evidence_collected": self.evidence_collected,
        }
    
    def _inst_to_dict(self, inst: EC2Instance) -> Dict:
        return {
            "instance_id": inst.instance_id, "state": inst.state.value,
            "is_compromised": inst.is_compromised, "has_snapshot": inst.has_forensic_snapshot
        }
    
    def _role_to_dict(self, role: IAMRole) -> Dict:
        return {
            "role_name": role.role_name, "is_compromised": role.is_compromised,
            "is_detached": role.is_detached, "has_backdoor": role.has_backdoor
        }
    
    def _bucket_to_dict(self, bucket: S3Bucket) -> Dict:
        return {
            "bucket_name": bucket.bucket_name, "is_public": bucket.is_public,
            "public_access_blocked": bucket.public_access_blocked
        }


# =============================================================================
# SCENARIO DEFINITIONS (Multi-Task Campaign - Mechanic #11)
# =============================================================================

SCENARIOS = {
    "easy": {
        "name": "Leaky S3 Bucket Discovery",
        "description": "Identify and secure a publicly exposed S3 bucket containing sensitive data",
        "difficulty_multiplier": 1.0,
        "max_steps": 15,
        "ground_truth_timeline": [
            "Public S3 bucket detected",
            "Credentials discovered in bucket",
            "Public access blocked"
        ],
        "required_flags": {"s3_public_identified", "credentials_found", "public_access_blocked"},
        "phase_weights": {"investigation": 0.5, "containment": 0.3, "eradication": 0.1, "recovery": 0.1},
        "mitre_techniques": ["T1530"],  # Data from Cloud Storage Object
        "hints": [
            "Start by reviewing SOC alerts",
            "Check S3 bucket policies for public access",
            "Look for sensitive files in public buckets"
        ]
    },
    "medium": {
        "name": "Credential Compromise Response",
        "description": "Leaked credentials from S3 were used; trace and revoke all compromised access",
        "difficulty_multiplier": 1.5,
        "max_steps": 25,
        "ground_truth_timeline": [
            "Leaked credentials identified",
            "Compromised IAM role discovered",
            "Role detached from resources",
            "Credentials rotated"
        ],
        "required_flags": {"leaked_creds_identified", "compromised_role_found", "role_detached", "credentials_rotated"},
        "phase_weights": {"investigation": 0.3, "containment": 0.4, "eradication": 0.2, "recovery": 0.1},
        "mitre_techniques": ["T1078", "T1552"],  # Valid Accounts, Credentials in Files
        "hints": [
            "Trace credential usage in CloudTrail",
            "Identify which roles were compromised",
            "Detach roles before rotating credentials"
        ]
    },
    "hard": {
        "name": "Full Incident Response - Ransomware",
        "description": "Compromised credentials led to IAM backdoor and ransomware deployment. Full IR required.",
        "difficulty_multiplier": 2.0,
        "max_steps": 40,
        "ground_truth_timeline": [
            "Ransomware indicators detected",
            "Backdoor IAM role identified",
            "Forensic snapshot taken",
            "Compromised instance isolated",
            "Backdoor role removed",
            "All credentials rotated",
            "Systems verified clean"
        ],
        "required_flags": {
            "ransomware_detected", "backdoor_identified", "forensic_snapshot_taken",
            "instance_isolated", "backdoor_removed", "all_creds_rotated", "systems_verified"
        },
        "phase_weights": {"investigation": 0.25, "containment": 0.25, "eradication": 0.3, "recovery": 0.2},
        "mitre_techniques": ["T1486", "T1098", "T1078"],  # Data Encrypted, Account Manipulation
        "hints": [
            "Preserve evidence with snapshots before isolation",
            "Look for recently created IAM roles with admin access",
            "Check for encryption activity on instances"
        ]
    }
}


# =============================================================================
# MAIN ENVIRONMENT CLASS
# =============================================================================

class CloudSOCEnv(gym.Env):
    """
    OpenEnv-CloudSOC Gymnasium Environment
    
    Implements all 12 required mechanics for LLM agent evaluation on
    cloud security incident response tasks.
    """
    
    metadata = {"render_modes": ["human", "json"]}
    
    def __init__(
        self,
        task: str = "easy",
        seed: Optional[int] = None,
        render_mode: str = "json",
        initial_state: Optional[CloudState] = None,
        verbose: bool = False
    ):
        """
        Initialize the CloudSOC environment.
        
        Args:
            task: Difficulty level ("easy", "medium", "hard")
            seed: Random seed for deterministic behavior (Mechanic #9)
            render_mode: Output format ("human" or "json")
            initial_state: Pre-existing state for multi-task campaigns (Mechanic #11)
            verbose: Enable detailed logging
        """
        super().__init__()
        
        self.task = task
        self.seed_value = seed
        self.render_mode = render_mode
        self.verbose = verbose
        self.scenario = SCENARIOS[task]
        
        # Deterministic seeding (Mechanic #9)
        if seed is not None:
            random.seed(seed)
            self._deterministic = True
        else:
            self._deterministic = False
        
        # Initialize or inherit state (Mechanic #11)
        if initial_state is not None:
            self.state = initial_state
            self._inherited_state = True
        else:
            self.state = CloudState()
            self._inherited_state = False
        
        # Environment tracking
        self.current_step = 0
        self.max_steps = self.scenario["max_steps"]
        self.total_reward = 0.0
        self.rewards_history: List[float] = []
        self.action_history: List[Dict] = []
        self.done = False
        self.last_action_error: Optional[str] = None
        
        # Phase scoring (Mechanic #8)
        self.phase_scores = {
            "investigation": 0.0,
            "containment": 0.0,
            "eradication": 0.0,
            "recovery": 0.0
        }
        
        # Tool usage analytics
        self.tool_usage: Dict[str, int] = {}
        self.query_costs = 0.0
        
        # Gymnasium spaces (simplified for LLM)
        self.action_space = spaces.Text(max_length=2048)
        self.observation_space = spaces.Text(max_length=8192)
        
        # Generate initial scenario
        self._generate_scenario()
    
    def _generate_scenario(self):
        """Generate the cloud environment state based on task difficulty"""
        
        if self._inherited_state and self.task != "easy":
            # Inherit and extend state for campaign continuity
            self._extend_scenario_state()
        else:
            # Fresh state generation
            self._generate_fresh_state()
        
        # Generate logs with noise (Mechanic #1)
        self._generate_logs()
        
        # Generate alerts
        self._generate_alerts()
    
    def _generate_fresh_state(self):
        """Generate brand new cloud state"""
        
        # Generate instance IDs deterministically if seeded
        instance_suffix = f"{random.randint(0, 0xffffffff):08x}"
        sg_suffix = f"{random.randint(0, 0xffffffff):08x}"
        
        # Create EC2 instances
        compromised_instance = EC2Instance(
            instance_id=f"i-{instance_suffix}",
            state=InstanceState.RUNNING,
            attached_roles=["WebServerRole"],
            security_groups=[f"sg-{sg_suffix}"],
            is_compromised=self.task in ["medium", "hard"],
            metadata={"name": "web-server-prod", "environment": "production"}
        )
        self.state.instances[compromised_instance.instance_id] = compromised_instance
        
        # Add benign instances (noise)
        for i in range(3):
            suffix = f"{random.randint(0, 0xffffffff):08x}"
            inst = EC2Instance(
                instance_id=f"i-{suffix}",
                state=InstanceState.RUNNING,
                attached_roles=[f"ServiceRole{i}"],
                security_groups=[f"sg-{random.randint(0, 0xffffffff):08x}"],
                is_compromised=False,
                metadata={"name": f"service-{i}", "environment": "production"}
            )
            self.state.instances[inst.instance_id] = inst
        
        # Create IAM roles
        web_role = IAMRole(
            role_name="WebServerRole",
            policies=["AmazonS3ReadOnly", "CloudWatchLogsFullAccess"],
            is_compromised=self.task in ["medium", "hard"]
        )
        self.state.roles["WebServerRole"] = web_role
        
        if self.task == "hard":
            backdoor_role = IAMRole(
                role_name="LegacyAdminBackup",
                policies=["AdministratorAccess"],
                is_compromised=True,
                has_backdoor=True
            )
            self.state.roles["LegacyAdminBackup"] = backdoor_role
        
        # Create S3 buckets
        leaky_bucket = S3Bucket(
            bucket_name="company-backup-2024",
            is_public=True,
            contains_credentials=True,
            leaked_keys=["AKIAIOSFODNN7EXAMPLE"]
        )
        self.state.buckets["company-backup-2024"] = leaky_bucket
        
        # Benign buckets
        for name in ["company-logs", "company-assets", "company-config"]:
            self.state.buckets[name] = S3Bucket(bucket_name=name, is_public=False)
        
        # Security groups
        self.state.security_groups[f"sg-{sg_suffix}"] = {
            "id": f"sg-{sg_suffix}",
            "name": "web-server-sg",
            "inbound_rules": [
                {"port": 443, "source": "0.0.0.0/0"},
                {"port": 22, "source": "10.0.0.0/8"}
            ],
            "is_isolated": False
        }
        
        # Leaked credentials
        self.state.credentials["AKIAIOSFODNN7EXAMPLE"] = {
            "key_id": "AKIAIOSFODNN7EXAMPLE",
            "user": "backup-service",
            "is_active": True,
            "is_compromised": True,
            "last_used": self._generate_timestamp(-2)
        }
        
        # Set ground truth timeline
        self.state.ground_truth_timeline = self.scenario["ground_truth_timeline"]
    
    def _extend_scenario_state(self):
        """Extend inherited state for harder scenarios"""
        
        if self.task == "medium":
            # Mark credentials as actively being abused
            for cred in self.state.credentials.values():
                if cred.get("is_compromised"):
                    cred["abuse_detected"] = True
                    cred["abuse_actions"] = ["AssumeRole", "ListBuckets", "GetObject"]
        
        elif self.task == "hard":
            # Add ransomware indicators
            for inst in self.state.instances.values():
                if inst.is_compromised:
                    inst.metadata["ransomware_indicators"] = True
                    inst.metadata["encrypted_files_detected"] = True
            
            # Add backdoor role if not exists
            if "LegacyAdminBackup" not in self.state.roles:
                self.state.roles["LegacyAdminBackup"] = IAMRole(
                    role_name="LegacyAdminBackup",
                    policies=["AdministratorAccess"],
                    is_compromised=True,
                    has_backdoor=True
                )
        
        # Update ground truth
        self.state.ground_truth_timeline = self.scenario["ground_truth_timeline"]
    
    def _generate_logs(self):
        """Generate logs with deceptive elements (Mechanic #1)"""
        
        base_time = datetime.utcnow() - timedelta(hours=24)
        
        # Attack indicators (true positives)
        attack_logs = [
            LogEntry(
                timestamp=self._generate_timestamp(-23),
                source="CloudTrail",
                message="AssumeRole called by AKIAIOSFODNN7EXAMPLE from IP 185.220.101.42",
                severity="WARN",
                is_attack_indicator=True
            ),
            LogEntry(
                timestamp=self._generate_timestamp(-22),
                source="S3",
                message="GetObject on s3://company-backup-2024/credentials.env from IP 185.220.101.42",
                severity="WARN",
                is_attack_indicator=True
            ),
            LogEntry(
                timestamp=self._generate_timestamp(-20),
                source="GuardDuty",
                message="UnauthorizedAccess:IAMUser/MaliciousIPCaller.Custom",
                severity="CRITICAL",
                is_attack_indicator=True
            ),
        ]
        
        if self.task == "hard":
            attack_logs.extend([
                LogEntry(
                    timestamp=self._generate_timestamp(-18),
                    source="CloudTrail",
                    message="CreateRole 'LegacyAdminBackup' with AdministratorAccess by compromised session",
                    severity="CRITICAL",
                    is_attack_indicator=True
                ),
                LogEntry(
                    timestamp=self._generate_timestamp(-15),
                    source="EC2",
                    message="Suspicious process execution: /tmp/.hidden/encrypt.sh on i-* instances",
                    severity="CRITICAL",
                    is_attack_indicator=True
                ),
            ])
        
        # Red herrings (Mechanic #1 - Deceptive Environment)
        red_herrings = [
            LogEntry(
                timestamp=self._generate_timestamp(-21),
                source="CloudTrail",
                message="Failed AssumeRole from IP 10.0.5.20 (security scanner - scheduled)",
                severity="WARN",
                is_red_herring=True
            ),
            LogEntry(
                timestamp=self._generate_timestamp(-19),
                source="IAM",
                message="Credential rotation reminder for service account 'deploy-bot'",
                severity="INFO",
                is_red_herring=True
            ),
            LogEntry(
                timestamp=self._generate_timestamp(-17),
                source="EC2",
                message="High CPU usage on i-abcd1234 (known batch job - analytics)",
                severity="WARN",
                is_red_herring=True
            ),
        ]
        
        # Noise (benign activity)
        noise_logs = []
        for i in range(15):
            noise_logs.append(LogEntry(
                timestamp=self._generate_timestamp(-random.randint(1, 23)),
                source=random.choice(["CloudTrail", "VPC", "Lambda", "ELB"]),
                message=random.choice([
                    f"Successful authentication from IP 10.0.{random.randint(0,255)}.{random.randint(0,255)}",
                    f"Lambda function 'data-processor-{random.randint(1,5)}' completed in {random.randint(100,5000)}ms",
                    f"Health check passed for target group tg-{random.randint(1000,9999)}",
                    f"VPC flow log: ACCEPT from 10.0.0.{random.randint(1,255)} to 10.0.1.{random.randint(1,255)}",
                    f"Auto-scaling event: DesiredCapacity changed from {random.randint(2,5)} to {random.randint(2,5)}",
                ]),
                severity="INFO",
                is_noise=True
            ))
        
        # Combine and shuffle
        all_logs = attack_logs + red_herrings + noise_logs
        random.shuffle(all_logs)
        self.state.logs = sorted(all_logs, key=lambda x: x.timestamp, reverse=True)
    
    def _generate_alerts(self):
        """Generate SOC alerts"""
        
        self.state.alerts = [
            Alert(
                alert_id="ALT-001",
                title="S3 Bucket Public Access Detected",
                severity="HIGH",
                source="AWS Config",
                timestamp=self._generate_timestamp(-24),
                details={"bucket": "company-backup-2024", "rule": "s3-bucket-public-read-prohibited"},
                is_true_positive=True
            ),
            Alert(
                alert_id="ALT-002",
                title="Unusual API Activity",
                severity="MEDIUM",
                source="GuardDuty",
                timestamp=self._generate_timestamp(-22),
                details={"finding_type": "UnauthorizedAccess:IAMUser/MaliciousIPCaller"},
                is_true_positive=True
            ),
            # False positive alert
            Alert(
                alert_id="ALT-003",
                title="Brute Force Attempt Detected",
                severity="LOW",
                source="GuardDuty",
                timestamp=self._generate_timestamp(-20),
                details={"source_ip": "10.0.5.20", "note": "Internal security scanner"},
                is_true_positive=False
            ),
        ]
        
        if self.task == "hard":
            self.state.alerts.append(Alert(
                alert_id="ALT-004",
                title="Ransomware Indicators Detected",
                severity="CRITICAL",
                source="CloudWatch",
                timestamp=self._generate_timestamp(-15),
                details={"indicators": ["mass file encryption", "ransom note creation"]},
                is_true_positive=True
            ))
    
    def _generate_timestamp(self, hours_offset: int) -> str:
        """Generate ISO timestamp with offset from now"""
        dt = datetime.utcnow() + timedelta(hours=hours_offset)
        return dt.strftime("%Y-%m-%dT%H:%M:%SZ")
    
    def reset(self, seed: Optional[int] = None, options: Optional[Dict] = None) -> Tuple[str, Dict]:
        """Reset environment to initial state"""
        
        if seed is not None:
            self.seed_value = seed
            random.seed(seed)
        
        self.current_step = 0
        self.total_reward = 0.0
        self.rewards_history = []
        self.action_history = []
        self.done = False
        self.last_action_error = None
        self.phase_scores = {k: 0.0 for k in self.phase_scores}
        self.tool_usage = {}
        self.query_costs = 0.0
        
        # Regenerate state
        if not self._inherited_state:
            self.state = CloudState()
            self._generate_scenario()
        else:
            # Reset investigation state but keep infrastructure
            self.state.discovered_flags = set()
            self.state.completed_actions = []
            self.state.agent_timeline = []
            self.state.incident_phase = IncidentPhase.INVESTIGATION
        
        observation = self._get_observation()
        info = self._get_info()
        
        return observation, info
    
    def step(self, action: str) -> Tuple[str, float, bool, bool, Dict]:
        """
        Execute one step in the environment.
        
        Args:
            action: JSON string containing tool call
            
        Returns:
            observation, reward, terminated, truncated, info
        """
        
        self.current_step += 1
        self.last_action_error = None
        reward = 0.0
        terminated = False
        truncated = False
        
        # Parse and validate action (Mechanic #7)
        try:
            action_dict = json.loads(action)
            tool_call = ToolCall(**action_dict)
        except json.JSONDecodeError as e:
            self.last_action_error = f"PARSE_ERROR: Invalid JSON - {str(e)}"
            reward = -0.02
            return self._finish_step(reward, terminated, truncated)
        except Exception as e:
            self.last_action_error = f"VALIDATION_ERROR: {str(e)}"
            reward = -0.02
            return self._finish_step(reward, terminated, truncated)
        
        # Track tool usage
        self.tool_usage[tool_call.tool] = self.tool_usage.get(tool_call.tool, 0) + 1
        
        # Record action
        self.action_history.append({
            "step": self.current_step,
            "tool": tool_call.tool,
            "args": tool_call.args,
            "thought": tool_call.thought
        })
        
        # Execute tool and get result
        result, step_reward, is_terminal, error = self._execute_tool(tool_call)
        
        reward = step_reward
        
        if error:
            self.last_action_error = error
        
        if is_terminal:
            terminated = True
        
        # Check for max steps (truncation)
        if self.current_step >= self.max_steps:
            truncated = True
        
        self.done = terminated or truncated
        
        return self._finish_step(reward, terminated, truncated)
    
    def _finish_step(self, reward: float, terminated: bool, truncated: bool) -> Tuple[str, float, bool, bool, Dict]:
        """Finalize step and return results"""
        
        self.total_reward += reward
        self.rewards_history.append(reward)
        
        observation = self._get_observation()
        info = self._get_info()
        
        return observation, reward, terminated, truncated, info
    
    def _execute_tool(self, tool_call: ToolCall) -> Tuple[Dict, float, bool, Optional[str]]:
        """
        Execute a tool call and return result.
        
        Returns:
            (result_dict, reward, is_terminal, error_message)
        """
        
        tool = tool_call.tool
        args = tool_call.args
        
        # Tool dispatch
        tool_handlers = {
            'aws.cloudwatch.query_basic': self._handle_cloudwatch_basic,
            'aws.cloudwatch.query_deep': self._handle_cloudwatch_deep,
            'aws.ec2.isolate': self._handle_ec2_isolate,
            'aws.ec2.terminate': self._handle_ec2_terminate,
            'aws.ec2.describe': self._handle_ec2_describe,
            'aws.ec2.snapshot': self._handle_ec2_snapshot,
            'aws.iam.detach_role': self._handle_iam_detach,
            'aws.iam.revoke_credentials': self._handle_iam_revoke,
            'aws.iam.describe_role': self._handle_iam_describe,
            'aws.iam.list_policies': self._handle_iam_list_policies,
            'aws.s3.get_bucket_policy': self._handle_s3_get_policy,
            'aws.s3.block_public_access': self._handle_s3_block_public,
            'aws.s3.list_objects': self._handle_s3_list_objects,
            'aws.rds.rotate_credentials': self._handle_rds_rotate,
            'aws.security_group.modify': self._handle_sg_modify,
            'aws.investigate': self._handle_investigate,
            'aws.soc.close_incident': self._handle_close_incident,
            'aws.soc.get_alerts': self._handle_get_alerts,
            'aws.guardduty.get_findings': self._handle_guardduty,
            'aws.cloudtrail.lookup_events': self._handle_cloudtrail,
            'aws.config.get_compliance': self._handle_config_compliance,
            'aws.ssm.run_command': self._handle_ssm_command,
            'aws.lambda.list_functions': self._handle_lambda_list,
            'aws.sts.get_caller_identity': self._handle_sts_identity,
        }
        
        handler = tool_handlers.get(tool)
        if not handler:
            return {}, -0.01, False, f"UNKNOWN_TOOL: {tool}"
        
        return handler(args)
    
    # =========================================================================
    # TOOL HANDLERS
    # =========================================================================
    
    def _handle_cloudwatch_basic(self, args: Dict) -> Tuple[Dict, float, bool, Optional[str]]:
        """Basic CloudWatch query - truncated results, low cost (Mechanic #2)"""
        
        cost = -0.01  # Resource cost
        self.query_costs += abs(cost)
        
        # Return only first 5 logs, summarized
        relevant_logs = [
            {"timestamp": log.timestamp, "source": log.source, "severity": log.severity}
            for log in self.state.logs[:5]
        ]
        
        # Discover basic flags
        new_flags = set()
        for log in self.state.logs[:5]:
            if log.is_attack_indicator and "S3" in log.source:
                new_flags.add("s3_activity_detected")
        
        flag_reward = self._process_discovered_flags(new_flags)
        
        return {
            "status": "success",
            "logs": relevant_logs,
            "total_count": len(self.state.logs),
            "note": "Truncated results. Use query_deep for full details."
        }, cost + flag_reward, False, None
    
    def _handle_cloudwatch_deep(self, args: Dict) -> Tuple[Dict, float, bool, Optional[str]]:
        """Deep CloudWatch query - full results, higher cost (Mechanic #2)"""
        
        cost = -0.05  # Higher resource cost
        self.query_costs += abs(cost)
        
        # Return full log details
        relevant_logs = [
            {
                "timestamp": log.timestamp,
                "source": log.source,
                "message": log.message,
                "severity": log.severity
            }
            for log in self.state.logs
            if not log.is_noise  # Filter pure noise but keep red herrings
        ]
        
        # Can discover more flags with deep query
        new_flags = set()
        for log in self.state.logs:
            if log.is_attack_indicator:
                if "S3" in log.source or "GetObject" in log.message:
                    new_flags.add("s3_public_identified")
                if "credentials" in log.message.lower():
                    new_flags.add("credentials_found")
                if "AssumeRole" in log.message and "185.220" in log.message:
                    new_flags.add("leaked_creds_identified")
                if "CreateRole" in log.message:
                    new_flags.add("backdoor_identified")
                if "ransomware" in log.message.lower() or "encrypt" in log.message.lower():
                    new_flags.add("ransomware_detected")
        
        flag_reward = self._process_discovered_flags(new_flags)
        self.phase_scores["investigation"] += 0.1 if new_flags else 0.0
        
        return {
            "status": "success",
            "logs": relevant_logs,
            "total_count": len(relevant_logs)
        }, cost + flag_reward, False, None
    
    def _handle_ec2_describe(self, args: Dict) -> Tuple[Dict, float, bool, Optional[str]]:
        """Describe EC2 instances"""
        
        instance_id = args.get("instance_id")
        
        if instance_id:
            inst = self.state.instances.get(instance_id)
            if not inst:
                return {"status": "error"}, -0.01, False, f"INSTANCE_NOT_FOUND: {instance_id}"
            instances = {instance_id: self._serialize_instance(inst)}
        else:
            instances = {
                iid: self._serialize_instance(inst) 
                for iid, inst in self.state.instances.items()
            }
        
        # Discover compromised instance flag
        new_flags = set()
        for iid, inst in self.state.instances.items():
            if inst.is_compromised and inst.metadata.get("ransomware_indicators"):
                new_flags.add("ransomware_detected")
        
        flag_reward = self._process_discovered_flags(new_flags)
        
        return {"status": "success", "instances": instances}, flag_reward, False, None
    
    def _serialize_instance(self, inst: EC2Instance) -> Dict:
        """Serialize instance to dict"""
        return {
            "instance_id": inst.instance_id,
            "state": inst.state.value,
            "attached_roles": inst.attached_roles,
            "security_groups": inst.security_groups,
            "has_forensic_snapshot": inst.has_forensic_snapshot,
            "metadata": inst.metadata
        }
    
    def _handle_ec2_snapshot(self, args: Dict) -> Tuple[Dict, float, bool, Optional[str]]:
        """Create forensic snapshot of instance"""
        
        instance_id = args.get("instance_id")
        if not instance_id:
            return {}, -0.01, False, "MISSING_PARAM: instance_id required"
        
        inst = self.state.instances.get(instance_id)
        if not inst:
            return {}, -0.01, False, f"INSTANCE_NOT_FOUND: {instance_id}"
        
        if inst.state == InstanceState.TERMINATED:
            return {}, -0.1, False, "ERROR: Cannot snapshot terminated instance"
        
        inst.has_forensic_snapshot = True
        self.state.completed_actions.append(f"snapshot:{instance_id}")
        
        new_flags = {"forensic_snapshot_taken"} if inst.is_compromised else set()
        flag_reward = self._process_discovered_flags(new_flags)
        
        self.phase_scores["containment"] += 0.15
        self.state.agent_timeline.append(f"Forensic snapshot taken for {instance_id}")
        
        return {
            "status": "success",
            "snapshot_id": f"snap-{hashlib.md5(instance_id.encode()).hexdigest()[:8]}",
            "instance_id": instance_id
        }, 0.05 + flag_reward, False, None
    
    def _handle_ec2_isolate(self, args: Dict) -> Tuple[Dict, float, bool, Optional[str]]:
        """Isolate EC2 instance - requires snapshot first (Mechanic #3)"""
        
        instance_id = args.get("instance_id")
        if not instance_id:
            return {}, -0.01, False, "MISSING_PARAM: instance_id required"
        
        inst = self.state.instances.get(instance_id)
        if not inst:
            return {}, -0.01, False, f"INSTANCE_NOT_FOUND: {instance_id}"
        
        # PRECONDITION CHECK (Mechanic #3)
        if not inst.has_forensic_snapshot:
            return {}, -0.05, False, "PRECONDITION_FAILED: Must take forensic snapshot before isolation"
        
        if inst.state == InstanceState.ISOLATED:
            return {"status": "already_isolated"}, 0.0, False, None
        
        inst.state = InstanceState.ISOLATED
        self.state.completed_actions.append(f"isolate:{instance_id}")
        
        new_flags = {"instance_isolated"} if inst.is_compromised else set()
        flag_reward = self._process_discovered_flags(new_flags)
        
        self.phase_scores["containment"] += 0.2
        self.state.agent_timeline.append(f"Instance {instance_id} isolated")
        
        return {
            "status": "success",
            "instance_id": instance_id,
            "new_state": "isolated"
        }, 0.1 + flag_reward, False, None
    
    def _handle_ec2_terminate(self, args: Dict) -> Tuple[Dict, float, bool, Optional[str]]:
        """ADVERSARIAL TRAP: Terminate instance destroys evidence (Mechanic #4)"""
        
        instance_id = args.get("instance_id")
        if not instance_id:
            return {}, -0.01, False, "MISSING_PARAM: instance_id required"
        
        inst = self.state.instances.get(instance_id)
        if not inst:
            return {}, -0.01, False, f"INSTANCE_NOT_FOUND: {instance_id}"
        
        if inst.is_compromised:
            # CATASTROPHIC ERROR - Evidence destruction
            inst.state = InstanceState.TERMINATED
            return {
                "status": "terminated",
                "warning": "CRITICAL: All forensic data has been permanently destroyed"
            }, -1.0, True, "CRITICAL: FORENSIC_DATA_LOSS_TERMINATION"
        
        # Non-compromised instance termination is just expensive
        inst.state = InstanceState.TERMINATED
        return {"status": "terminated"}, -0.1, False, None
    
    def _handle_iam_detach(self, args: Dict) -> Tuple[Dict, float, bool, Optional[str]]:
        """Detach IAM role from resources"""
        
        role_name = args.get("role_name")
        if not role_name:
            return {}, -0.01, False, "MISSING_PARAM: role_name required"
        
        role = self.state.roles.get(role_name)
        if not role:
            return {}, -0.01, False, f"ROLE_NOT_FOUND: {role_name}"
        
        if role.is_detached:
            return {"status": "already_detached"}, 0.0, False, None
        
        role.is_detached = True
        self.state.completed_actions.append(f"detach_role:{role_name}")
        
        new_flags = set()
        if role.is_compromised:
            new_flags.add("role_detached")
        if role.has_backdoor:
            new_flags.add("backdoor_removed")
        
        flag_reward = self._process_discovered_flags(new_flags)
        self.phase_scores["eradication"] += 0.15
        self.state.agent_timeline.append(f"Role {role_name} detached")
        
        return {
            "status": "success",
            "role_name": role_name,
            "detached": True
        }, 0.05 + flag_reward, False, None
    
    def _handle_iam_revoke(self, args: Dict) -> Tuple[Dict, float, bool, Optional[str]]:
        """Revoke IAM credentials"""
        
        key_id = args.get("key_id") or args.get("resource_id")
        if not key_id:
            return {}, -0.01, False, "MISSING_PARAM: key_id required"
        
        cred = self.state.credentials.get(key_id)
        if not cred:
            return {}, -0.01, False, f"CREDENTIAL_NOT_FOUND: {key_id}"
        
        cred["is_active"] = False
        cred["revoked_at"] = self._generate_timestamp(0)
        
        new_flags = {"credentials_rotated"} if cred.get("is_compromised") else set()
        flag_reward = self._process_discovered_flags(new_flags)
        
        self.phase_scores["eradication"] += 0.1
        self.state.agent_timeline.append(f"Credentials {key_id} revoked")
        
        return {
            "status": "success",
            "key_id": key_id,
            "revoked": True
        }, 0.05 + flag_reward, False, None
    
    def _handle_iam_describe(self, args: Dict) -> Tuple[Dict, float, bool, Optional[str]]:
        """Describe IAM role"""
        
        role_name = args.get("role_name")
        if not role_name:
            # List all roles
            roles = {name: self._serialize_role(role) for name, role in self.state.roles.items()}
            return {"status": "success", "roles": roles}, 0.0, False, None
        
        role = self.state.roles.get(role_name)
        if not role:
            return {}, -0.01, False, f"ROLE_NOT_FOUND: {role_name}"
        
        new_flags = set()
        if role.is_compromised:
            new_flags.add("compromised_role_found")
        if role.has_backdoor:
            new_flags.add("backdoor_identified")
        
        flag_reward = self._process_discovered_flags(new_flags)
        
        return {
            "status": "success",
            "role": self._serialize_role(role)
        }, flag_reward, False, None
    
    def _serialize_role(self, role: IAMRole) -> Dict:
        return {
            "role_name": role.role_name,
            "policies": role.policies,
            "is_detached": role.is_detached,
            "credentials_rotated": role.credentials_rotated
        }
    
    def _handle_iam_list_policies(self, args: Dict) -> Tuple[Dict, float, bool, Optional[str]]:
        """List IAM policies"""
        all_policies = set()
        for role in self.state.roles.values():
            all_policies.update(role.policies)
        return {"status": "success", "policies": list(all_policies)}, 0.0, False, None
    
    def _handle_s3_get_policy(self, args: Dict) -> Tuple[Dict, float, bool, Optional[str]]:
        """Get S3 bucket policy"""
        
        bucket_name = args.get("bucket_name")
        if not bucket_name:
            return {}, -0.01, False, "MISSING_PARAM: bucket_name required"
        
        bucket = self.state.buckets.get(bucket_name)
        if not bucket:
            return {}, -0.01, False, f"BUCKET_NOT_FOUND: {bucket_name}"
        
        new_flags = set()
        if bucket.is_public:
            new_flags.add("s3_public_identified")
        
        flag_reward = self._process_discovered_flags(new_flags)
        self.phase_scores["investigation"] += 0.05 if new_flags else 0.0
        
        return {
            "status": "success",
            "bucket_name": bucket_name,
            "is_public": bucket.is_public,
            "public_access_blocked": bucket.public_access_blocked,
            "policy": {"Effect": "Allow", "Principal": "*"} if bucket.is_public else {}
        }, flag_reward, False, None
    
    def _handle_s3_block_public(self, args: Dict) -> Tuple[Dict, float, bool, Optional[str]]:
        """Block public access to S3 bucket"""
        
        bucket_name = args.get("bucket_name")
        if not bucket_name:
            return {}, -0.01, False, "MISSING_PARAM: bucket_name required"
        
        bucket = self.state.buckets.get(bucket_name)
        if not bucket:
            return {}, -0.01, False, f"BUCKET_NOT_FOUND: {bucket_name}"
        
        if bucket.public_access_blocked:
            return {"status": "already_blocked"}, 0.0, False, None
        
        bucket.public_access_blocked = True
        bucket.is_public = False
        self.state.completed_actions.append(f"block_public:{bucket_name}")
        
        new_flags = {"public_access_blocked"} if bucket.contains_credentials else set()
        flag_reward = self._process_discovered_flags(new_flags)
        
        self.phase_scores["containment"] += 0.15
        self.state.agent_timeline.append(f"Public access blocked for {bucket_name}")
        
        return {
            "status": "success",
            "bucket_name": bucket_name,
            "public_access_blocked": True
        }, 0.1 + flag_reward, False, None
    
    def _handle_s3_list_objects(self, args: Dict) -> Tuple[Dict, float, bool, Optional[str]]:
        """List S3 bucket objects"""
        
        bucket_name = args.get("bucket_name")
        if not bucket_name:
            return {}, -0.01, False, "MISSING_PARAM: bucket_name required"
        
        bucket = self.state.buckets.get(bucket_name)
        if not bucket:
            return {}, -0.01, False, f"BUCKET_NOT_FOUND: {bucket_name}"
        
        objects = ["backup-2024-01.tar.gz", "config.yaml", "logs/"]
        if bucket.contains_credentials:
            objects.append("credentials.env")
        
        new_flags = set()
        if bucket.contains_credentials:
            new_flags.add("credentials_found")
        
        flag_reward = self._process_discovered_flags(new_flags)
        
        return {
            "status": "success",
            "bucket_name": bucket_name,
            "objects": objects
        }, flag_reward, False, None
    
    def _handle_rds_rotate(self, args: Dict) -> Tuple[Dict, float, bool, Optional[str]]:
        """Rotate RDS credentials - requires role detachment first (Mechanic #3)"""
        
        resource_id = args.get("resource_id")
        if not resource_id:
            return {}, -0.01, False, "MISSING_PARAM: resource_id required"
        
        # PRECONDITION: Check if compromised roles are detached
        compromised_attached = any(
            role.is_compromised and not role.is_detached
            for role in self.state.roles.values()
        )
        
        if compromised_attached:
            return {}, -0.05, False, "PRECONDITION_FAILED: Detach compromised IAM roles before rotating credentials"
        
        self.state.completed_actions.append(f"rotate_creds:{resource_id}")
        
        # Check if all credentials are now rotated
        all_rotated = all(
            not cred.get("is_active", True) 
            for cred in self.state.credentials.values() 
            if cred.get("is_compromised")
        )
        
        new_flags = {"all_creds_rotated"} if all_rotated else set()
        flag_reward = self._process_discovered_flags(new_flags)
        
        self.phase_scores["recovery"] += 0.15
        self.state.agent_timeline.append(f"Credentials rotated for {resource_id}")
        
        return {
            "status": "success",
            "resource_id": resource_id,
            "credentials_rotated": True
        }, 0.05 + flag_reward, False, None
    
    def _handle_sg_modify(self, args: Dict) -> Tuple[Dict, float, bool, Optional[str]]:
        """Modify security group"""
        
        sg_id = args.get("security_group_id")
        action = args.get("action")
        
        if not sg_id or not action:
            return {}, -0.01, False, "MISSING_PARAM: security_group_id and action required"
        
        sg = self.state.security_groups.get(sg_id)
        if not sg:
            return {}, -0.01, False, f"SECURITY_GROUP_NOT_FOUND: {sg_id}"
        
        if action == "isolate":
            sg["is_isolated"] = True
            sg["inbound_rules"] = []
            self.phase_scores["containment"] += 0.1
        elif action == "restore":
            sg["is_isolated"] = False
            self.phase_scores["recovery"] += 0.05
        
        return {"status": "success", "security_group_id": sg_id, "action": action}, 0.02, False, None
    
    def _handle_investigate(self, args: Dict) -> Tuple[Dict, float, bool, Optional[str]]:
        """General investigation tool"""
        
        target_type = args.get("target_type")
        target_id = args.get("target_id")
        
        if not target_type or not target_id:
            return {}, -0.01, False, "MISSING_PARAM: target_type and target_id required"
        
        result = {"status": "success", "target_type": target_type, "target_id": target_id}
        new_flags = set()
        
        if target_type == "ip" and "185.220" in target_id:
            result["findings"] = {
                "reputation": "malicious",
                "associated_campaigns": ["Credential Theft", "Cloud Exploitation"],
                "geo": "TOR Exit Node"
            }
            new_flags.add("malicious_ip_confirmed")
        elif target_type == "user":
            result["findings"] = {"last_activity": self._generate_timestamp(-1)}
        elif target_type == "role":
            role = self.state.roles.get(target_id)
            if role:
                result["findings"] = self._serialize_role(role)
                if role.is_compromised:
                    new_flags.add("compromised_role_found")
        
        flag_reward = self._process_discovered_flags(new_flags)
        self.phase_scores["investigation"] += 0.05 if new_flags else 0.0
        
        return result, flag_reward, False, None
    
    def _handle_get_alerts(self, args: Dict) -> Tuple[Dict, float, bool, Optional[str]]:
        """Get SOC alerts"""
        
        alerts = [
            {
                "alert_id": a.alert_id,
                "title": a.title,
                "severity": a.severity,
                "source": a.source,
                "timestamp": a.timestamp,
                "details": a.details
            }
            for a in self.state.alerts
        ]
        
        return {"status": "success", "alerts": alerts, "count": len(alerts)}, 0.0, False, None
    
    def _handle_guardduty(self, args: Dict) -> Tuple[Dict, float, bool, Optional[str]]:
        """Get GuardDuty findings"""
        
        findings = [
            {
                "id": f"gd-{i}",
                "type": log.message.split(":")[0] if ":" in log.message else "UnauthorizedAccess",
                "severity": log.severity,
                "description": log.message
            }
            for i, log in enumerate(self.state.logs)
            if log.source == "GuardDuty" and log.is_attack_indicator
        ]
        
        new_flags = set()
        if findings:
            new_flags.add("guardduty_reviewed")
        
        flag_reward = self._process_discovered_flags(new_flags)
        
        return {"status": "success", "findings": findings}, flag_reward, False, None
    
    def _handle_cloudtrail(self, args: Dict) -> Tuple[Dict, float, bool, Optional[str]]:
        """Lookup CloudTrail events"""
        
        events = [
            {
                "timestamp": log.timestamp,
                "event_name": log.message.split()[0] if log.message else "Unknown",
                "source_ip": random.choice(MALICIOUS_IPS) if log.is_attack_indicator else random.choice(BENIGN_IPS),
                "user": "compromised-user" if log.is_attack_indicator else "system",
                "details": log.message
            }
            for log in self.state.logs
            if log.source == "CloudTrail"
        ]
        
        return {"status": "success", "events": events}, 0.0, False, None
    
    def _handle_config_compliance(self, args: Dict) -> Tuple[Dict, float, bool, Optional[str]]:
        """Get AWS Config compliance status"""
        
        rules = [
            {"rule": "s3-bucket-public-read-prohibited", "compliance": "NON_COMPLIANT" if any(b.is_public for b in self.state.buckets.values()) else "COMPLIANT"},
            {"rule": "iam-root-access-key-check", "compliance": "COMPLIANT"},
            {"rule": "ec2-instance-no-public-ip", "compliance": "COMPLIANT"},
            {"rule": "encrypted-volumes", "compliance": "COMPLIANT"},
        ]
        
        new_flags = set()
        if any(r["compliance"] == "NON_COMPLIANT" for r in rules):
            new_flags.add("compliance_issue_found")
        
        flag_reward = self._process_discovered_flags(new_flags)
        
        return {"status": "success", "compliance_rules": rules}, flag_reward, False, None
    
    def _handle_ssm_command(self, args: Dict) -> Tuple[Dict, float, bool, Optional[str]]:
        """Run SSM command on instance (for verification)"""
        
        instance_id = args.get("instance_id")
        command = args.get("command", "")
        
        if not instance_id:
            return {}, -0.01, False, "MISSING_PARAM: instance_id required"
        
        inst = self.state.instances.get(instance_id)
        if not inst:
            return {}, -0.01, False, f"INSTANCE_NOT_FOUND: {instance_id}"
        
        if inst.state != InstanceState.RUNNING:
            return {}, -0.01, False, f"INSTANCE_NOT_RUNNING: {instance_id} is {inst.state.value}"
        
        # Simulate command output
        if "check" in command.lower() or "verify" in command.lower():
            if inst.is_compromised and inst.metadata.get("ransomware_indicators"):
                output = "WARNING: Suspicious processes detected\n/tmp/.hidden/encrypt.sh (running)"
            else:
                output = "System check passed. No anomalies detected."
            
            new_flags = {"systems_verified"} if not inst.is_compromised else set()
            flag_reward = self._process_discovered_flags(new_flags)
            self.phase_scores["recovery"] += 0.1
            
            return {"status": "success", "output": output}, flag_reward, False, None
        
        return {"status": "success", "output": "Command executed"}, 0.0, False, None
    
    def _handle_lambda_list(self, args: Dict) -> Tuple[Dict, float, bool, Optional[str]]:
        """List Lambda functions"""
        
        functions = [
            {"name": "data-processor", "runtime": "python3.9", "last_modified": self._generate_timestamp(-48)},
            {"name": "api-handler", "runtime": "nodejs18.x", "last_modified": self._generate_timestamp(-72)},
            {"name": "auth-validator", "runtime": "python3.9", "last_modified": self._generate_timestamp(-24)},
        ]
        
        return {"status": "success", "functions": functions}, 0.0, False, None
    
    def _handle_sts_identity(self, args: Dict) -> Tuple[Dict, float, bool, Optional[str]]:
        """Get current caller identity"""
        
        return {
            "status": "success",
            "account": "123456789012",
            "arn": "arn:aws:iam::123456789012:user/security-analyst",
            "user_id": "AIDAEXAMPLEID"
        }, 0.0, False, None
    
    def _handle_close_incident(self, args: Dict) -> Tuple[Dict, float, bool, Optional[str]]:
        """
        Close incident with timeline (Mechanic #12)
        
        This is the FINAL action to complete a task successfully.
        """
        
        timeline = args.get("timeline", [])
        incident_id = args.get("incident_id", "INC-001")
        
        if not timeline:
            return {}, -0.1, False, "MISSING_PARAM: timeline array required"
        
        if not isinstance(timeline, list) or len(timeline) < 2:
            return {}, -0.1, False, "INVALID_PARAM: timeline must be an array with at least 2 events"
        
        # Store agent's timeline
        self.state.agent_timeline = timeline
        
        # Grade the timeline (Mechanic #12)
        timeline_score = self._grade_timeline(timeline)
        
        # Check if required flags are discovered
        required_flags = self.scenario["required_flags"]
        discovered = self.state.discovered_flags
        flags_complete = required_flags.issubset(discovered)
        
        # Calculate final scores
        final_scores = self.calculate_final_score()
        
        # Determine success
        success = flags_complete and timeline_score >= 0.5
        
        self.state.incident_phase = IncidentPhase.CLOSED
        
        return {
            "status": "incident_closed",
            "incident_id": incident_id,
            "timeline_submitted": timeline,
            "timeline_score": round(timeline_score, 2),
            "flags_complete": flags_complete,
            "discovered_flags": list(discovered),
            "required_flags": list(required_flags),
            "final_scores": final_scores,
            "success": success
        }, timeline_score * 0.5, True, None
    
    def _grade_timeline(self, agent_timeline: List[str]) -> float:
        """
        Grade the agent's incident timeline against ground truth.
        
        Uses multi-factor scoring:
        1. Keyword matching for event detection
        2. Order preservation bonus
        3. Completeness bonus
        4. Penalty for extraneous/wrong events
        """
        
        ground_truth = self.state.ground_truth_timeline
        if not ground_truth:
            return 0.5  # Default score if no ground truth
        
        if not agent_timeline:
            return 0.0
        
        score = 0.0
        max_score = len(ground_truth)
        matched_indices = []
        
        agent_lower = [t.lower() for t in agent_timeline]
        
        # Score each ground truth event
        for truth_idx, truth_event in enumerate(ground_truth):
            truth_lower = truth_event.lower()
            truth_words = set(truth_lower.split())
            
            best_match_score = 0.0
            best_match_idx = -1
            
            for agent_idx, agent_event in enumerate(agent_lower):
                agent_words = set(agent_event.split())
                
                # Jaccard similarity
                intersection = len(truth_words & agent_words)
                union = len(truth_words | agent_words)
                similarity = intersection / union if union > 0 else 0
                
                # Keyword bonus for critical terms
                critical_keywords = {"snapshot", "isolate", "detach", "rotate", "blocked", "revoke", "backdoor", "ransomware", "credential"}
                keyword_matches = sum(1 for kw in critical_keywords if kw in agent_event)
                keyword_bonus = min(0.2, keyword_matches * 0.1)
                
                match_score = similarity + keyword_bonus
                
                if match_score > best_match_score:
                    best_match_score = match_score
                    best_match_idx = agent_idx
            
            if best_match_score >= 0.3:  # Threshold for acceptable match
                score += best_match_score
                matched_indices.append((truth_idx, best_match_idx))
        
        # Order preservation bonus (check if matched events maintain relative order)
        if len(matched_indices) >= 2:
            order_preserved = all(
                matched_indices[i][1] < matched_indices[i+1][1]
                for i in range(len(matched_indices) - 1)
                if matched_indices[i][1] != -1 and matched_indices[i+1][1] != -1
            )
            if order_preserved:
                score += 0.5  # Bonus for correct ordering
        
        # Completeness bonus
        coverage = len([m for m in matched_indices if m[1] != -1]) / max_score
        score += coverage * 0.3
        
        # Normalize to 0-1 range
        normalized = min(1.0, score / (max_score + 0.8))
        
        # Scale to 0.1 - 0.9 per requirements
        return 0.1 + (normalized * 0.8)
    
    def _process_discovered_flags(self, new_flags: set) -> float:
        """
        Process newly discovered flags and return reward.
        Implements Mechanic #5 (Gradient Scoring).
        """
        
        actually_new = new_flags - self.state.discovered_flags
        self.state.discovered_flags.update(actually_new)
        
        # Reward per new flag (Mechanic #5)
        reward = 0.02 * len(actually_new)
        
        return reward
    
    def calculate_final_score(self) -> Dict[str, float]:
        """
        Calculate final score breakdown (Mechanic #8).
        
        Returns structured scoring across IR phases.
        """
        
        weights = self.scenario["phase_weights"]
        
        # Normalize phase scores to 0-1 range
        max_scores = {
            "investigation": 0.5,
            "containment": 0.6,
            "eradication": 0.4,
            "recovery": 0.3
        }
        
        normalized = {}
        for phase, score in self.phase_scores.items():
            max_s = max_scores.get(phase, 0.5)
            normalized[phase] = min(1.0, score / max_s) if max_s > 0 else 0.0
        
        # Weighted average
        weighted_total = sum(
            normalized[phase] * weights[phase]
            for phase in weights
        )
        
        # Adjust weighted_total to range 0.1 to 0.9 rigidly
        weighted_total = 0.1 + (weighted_total * 0.8)
        
        return {
            "investigation": round(0.1 + (normalized["investigation"] * 0.8), 3),
            "containment": round(0.1 + (normalized["containment"] * 0.8), 3),
            "eradication": round(0.1 + (normalized["eradication"] * 0.8), 3),
            "recovery": round(0.1 + (normalized["recovery"] * 0.8), 3),
            "weighted_total": round(weighted_total, 3),
            "query_costs": round(self.query_costs, 3),
            "total_reward": round(self.total_reward, 3)
        }
    
    def _get_observation(self) -> str:
        """Get current observation as JSON string"""
        
        obs = {
            "step": self.current_step,
            "phase": self.state.incident_phase.value,
            "discovered_flags": list(self.state.discovered_flags),
            "last_error": self.last_action_error,
            "alerts_pending": len([a for a in self.state.alerts if a.is_true_positive]),
            "instances_count": len(self.state.instances),
            "compromised_instances": len([i for i in self.state.instances.values() if i.is_compromised]),
        }
        
        if self.action_history:
            last_action = self.action_history[-1]
            obs["last_action"] = {
                "tool": last_action["tool"],
                "result": "error" if self.last_action_error else "success"
            }
        
        return json.dumps(obs, indent=2)
    
    def _get_info(self) -> Dict:
        """Get additional info dict"""
        return {
            "task": self.task,
            "scenario": self.scenario["name"],
            "current_step": self.current_step,
            "max_steps": self.max_steps,
            "total_reward": round(self.total_reward, 4),
            "discovered_flags": list(self.state.discovered_flags),
            "required_flags": list(self.scenario["required_flags"]),
            "phase_scores": self.phase_scores.copy(),
            "tool_usage": self.tool_usage.copy(),
            "last_action_error": self.last_action_error
        }
    
    def render(self, mode: Optional[str] = None) -> Optional[str]:
        """Render current state"""
        
        mode = mode or self.render_mode
        
        if mode == "json":
            return json.dumps({
                "task": self.task,
                "step": self.current_step,
                "state": {
                    "instances": {k: self._serialize_instance(v) for k, v in self.state.instances.items()},
                    "discovered_flags": list(self.state.discovered_flags),
                    "phase": self.state.incident_phase.value
                },
                "scores": self.calculate_final_score()
            }, indent=2)
        
        elif mode == "human":
            lines = [
                f"=== CloudSOC Environment ===",
                f"Task: {self.scenario['name']} ({self.task})",
                f"Step: {self.current_step}/{self.max_steps}",
                f"Phase: {self.state.incident_phase.value}",
                f"Flags: {len(self.state.discovered_flags)}/{len(self.scenario['required_flags'])}",
                f"Total Reward: {self.total_reward:.4f}",
                f"Last Error: {self.last_action_error or 'None'}"
            ]
            return "\n".join(lines)
        
        return None
    
    def close(self):
        """Clean up environment"""
        pass
    
    def get_state_for_next_task(self) -> CloudState:
        """
        Export current state for multi-task campaign (Mechanic #11).
        
        Returns the state object to pass to the next task's environment.
        """
        return self.state
    
    def get_system_prompt(self) -> str:
        """
        Generate the system prompt for the LLM agent.
        Implements Mechanic #10 (Chain-of-Thought Prompting).
        """
        
        available_tools = [
            "aws.cloudwatch.query_basic - Query logs (truncated, cost: -0.01)",
            "aws.cloudwatch.query_deep - Query logs (full, cost: -0.05)",
            "aws.ec2.describe - Describe EC2 instances",
            "aws.ec2.snapshot - Create forensic snapshot (REQUIRED before isolation)",
            "aws.ec2.isolate - Isolate compromised instance",
            "aws.ec2.terminate - ⚠️ DANGER: Permanently terminates instance, destroys evidence",
            "aws.iam.describe_role - Describe IAM role",
            "aws.iam.detach_role - Detach IAM role from resources",
            "aws.iam.revoke_credentials - Revoke compromised credentials",
            "aws.iam.list_policies - List IAM policies",
            "aws.s3.get_bucket_policy - Check S3 bucket policy",
            "aws.s3.block_public_access - Block public access to bucket",
            "aws.s3.list_objects - List bucket contents",
            "aws.rds.rotate_credentials - Rotate DB credentials (requires role detachment first)",
            "aws.security_group.modify - Modify security group (isolate/restore)",
            "aws.investigate - Investigate IP, user, role, or instance",
            "aws.soc.get_alerts - Get current SOC alerts",
            "aws.guardduty.get_findings - Get GuardDuty findings",
            "aws.cloudtrail.lookup_events - Lookup CloudTrail events",
            "aws.config.get_compliance - Get AWS Config compliance status",
            "aws.ssm.run_command - Run command on instance for verification",
            "aws.lambda.list_functions - List Lambda functions",
            "aws.sts.get_caller_identity - Get current caller identity",
            "aws.soc.close_incident - Close incident with timeline (FINAL action)"
        ]
        
        # Get progress hints based on current state
        progress_hints = []
        discovered = self.state.discovered_flags
        required = self.scenario["required_flags"]
        remaining = required - discovered
        
        if len(discovered) == 0:
            progress_hints.append("Start by gathering information: check alerts and logs")
        elif len(remaining) <= 2:
            progress_hints.append("You're close! Complete remaining actions and close the incident")
        
        return f"""You are an expert Cloud Security Analyst responding to a security incident in AWS.

## INCIDENT: {self.scenario['name']}
{self.scenario['description']}

## RESPONSE FORMAT (MANDATORY)
You MUST respond with valid JSON in this exact format:
```json
{{
  "thought": "Your 1-sentence reasoning about what to do next and why",
  "tool": "tool_name",
  "args": {{"param": "value"}}
}}
```

## ⚠️ CRITICAL WARNINGS - READ CAREFULLY
1. **NEVER use aws.ec2.terminate on compromised instances** - This permanently destroys all forensic evidence and fails the mission
2. **ALWAYS take a forensic snapshot BEFORE isolating** - Isolation without snapshot = precondition failure
3. **ALWAYS detach compromised IAM roles BEFORE rotating credentials** - Rotation without detachment = precondition failure
4. **Query costs matter**: basic=-0.01, deep=-0.05 - Balance thoroughness with efficiency

## INCIDENT RESPONSE PROCEDURE
1. **Investigation**: Gather information (alerts, logs, findings)
2. **Containment**: Stop the spread (snapshot → isolate, block access)
3. **Eradication**: Remove threat (detach roles, revoke credentials)
4. **Recovery**: Restore operations (rotate credentials, verify systems)
5. **Close**: Submit timeline with aws.soc.close_incident

## AVAILABLE TOOLS
{chr(10).join(f'• {t}' for t in available_tools)}

## TOOL PARAMETER SCHEMAS (MANDATORY - Use exact parameter names)
• aws.s3.get_bucket_policy: args: {{"bucket_name": "string"}}
• aws.s3.list_objects: args: {{"bucket_name": "string"}}
• aws.s3.block_public_access: args: {{"bucket_name": "string"}}
• aws.ec2.describe: args: {{"instance_id": "string"}} (optional)
• aws.ec2.isolate: args: {{"instance_id": "string"}}
• aws.ec2.snapshot: args: {{"instance_id": "string"}}
• aws.ec2.terminate: args: {{"instance_id": "string"}}
• aws.iam.describe_role: args: {{"role_name": "string"}}
• aws.iam.detach_role: args: {{"role_name": "string", "resource_id": "string"}}
• aws.iam.revoke_credentials: args: {{"principal": "string", "credential_type": "string"}}
• aws.cloudwatch.query_basic: args: {{"log_group": "string", "query": "string"}}
• aws.cloudwatch.query_deep: args: {{"log_group": "string", "query": "string"}}
• aws.cloudtrail.lookup_events: args: {{"event_names": [list], "max_results": int}}
• aws.soc.close_incident: args: {{"timeline": [list of strings]}}
• aws.investigate: args: {{"query": "string"}}
• aws.security_group.modify: args: {{"group_id": "string", "action": "string"}}
• aws.rds.rotate_credentials: args: {{"instance_id": "string"}}
• aws.guardduty.get_findings: args: {{}} (no args required)
• aws.soc.get_alerts: args: {{}} (no args required)
• aws.config.get_compliance: args: {{}} (no args required)
• aws.ssm.run_command: args: {{"instance_id": "string", "command": "string"}}
• aws.lambda.list_functions: args: {{}} (no args required)
• aws.sts.get_caller_identity: args: {{}} (no args required)

## CURRENT STATUS
• Step: {self.current_step}/{self.max_steps}
• Phase: {self.state.incident_phase.value}
• Progress: {len(self.state.discovered_flags)}/{len(self.scenario['required_flags'])} flags discovered
• Query costs so far: {self.query_costs:.2f}
{chr(10).join(f'• Hint: {h}' for h in progress_hints)}

## FINAL STEP
To complete this incident, call `aws.soc.close_incident` with a timeline array summarizing the incident chronologically. Example:
```json
{{"thought": "Investigation complete, closing incident", "tool": "aws.soc.close_incident", "args": {{"timeline": ["S3 bucket found public", "Credentials discovered", "Access blocked"]}}}}
```

Think step-by-step. Each action should move you closer to resolution."""

    
# =============================================================================
# HELPER FUNCTIONS
# =============================================================================

def create_campaign() -> List[Tuple[str, Optional[CloudState]]]:
    """
    Create a multi-task campaign (Mechanic #11).
    
    Returns list of (task_name, initial_state) tuples.
    """
    return [
        ("easy", None),
        ("medium", None),  # Will receive state from easy
        ("hard", None),    # Will receive state from medium
    ]


def run_campaign_task(
    task: str,
    initial_state: Optional[CloudState] = None,
    seed: Optional[int] = None
) -> Tuple[CloudSOCEnv, CloudState]:
    """
    Run a single task in the campaign.
    
    Returns (env, final_state) for chaining to next task.
    """
    env = CloudSOCEnv(task=task, initial_state=initial_state, seed=seed)
    return env, env.state


# =============================================================================
# MAIN (for testing)
# =============================================================================

if __name__ == "__main__":
    # Test environment creation
    env = CloudSOCEnv(task="easy", seed=42)
    obs, info = env.reset()
    
    print("=== CloudSOC Environment Test ===")
    print(f"Task: {env.task}")
    print(f"Scenario: {env.scenario['name']}")
    print(f"Max Steps: {env.max_steps}")
    print(f"\nInitial Observation:\n{obs}")
    print(f"\nSystem Prompt:\n{env.get_system_prompt()[:500]}...")
    
    # Test a simple action
    test_action = json.dumps({
        "thought": "First, I need to check the SOC alerts to understand the incident",
        "tool": "aws.soc.get_alerts",
        "args": {}
    })
    
    obs, reward, term, trunc, info = env.step(test_action)
    print(f"\n=== After Step 1 ===")
    print(f"Reward: {reward}")
    print(f"Observation:\n{obs}")
    
    env.close()