bothari01/secops-env
0
1"""2Tool Simulator - Mock AWS CLI Execution.3 4Simulates command execution without external dependencies.5Allows agents to execute simulated AWS commands and receive realistic responses.6"""7 8import re9import time10from typing import Any, Dict, List, Optional, Tuple11from dataclasses import dataclass, field12from datetime import datetime13 14 15@dataclass16class CommandResult:17 success: bool18 output: str19 error: Optional[str] = None20 command: str = ""21 timestamp: datetime = field(default_factory=datetime.now)22 23 24@dataclass25class CloudResource:26 resource_type: str27 resource_id: str28 properties: Dict[str, Any] = field(default_factory=dict)29 public: bool = False30 disabled: bool = False31 32 33class ToolSimulator:34 """35 Simulates command execution without external dependencies.36 37 Supports simulated AWS CLI commands for:38 - S3 bucket operations39 - IAM user management40 - EC2 instance queries41 - Security group configurations42 """43 44 def __init__(self):45 self.execution_log: List[Dict[str, Any]] = []46 self.cloud_state: Dict[str, Dict[str, CloudResource]] = {47 "s3_buckets": {},48 "iam_users": {},49 "ec2_instances": {},50 "security_groups": {},51 }52 self._execution_count = 053 54 def execute_aws_command(55 self, command: str, args: Dict[str, Any] = None56 ) -> CommandResult:57 """58 Parse and execute simulated AWS CLI commands.59 60 Args:61 command: The AWS CLI command string62 args: Optional parsed arguments63 64 Returns:65 CommandResult with success status, output, and optional error66 """67 self._execution_count += 168 args = args or {}69 70 if "s3api" in command or "s3" in command.lower():71 return self._execute_s3_command(command, args)72 elif "iam" in command.lower():73 return self._execute_iam_command(command, args)74 elif "ec2" in command.lower():75 return self._execute_ec2_command(command, args)76 elif "describe-security-groups" in command:77 return self._execute_security_group_command(command, args)78 else:79 return CommandResult(80 success=False,81 output="",82 error=f"Unknown command type: {command}",83 command=command,84 )85 86 def _execute_s3_command(self, command: str, args: Dict[str, Any]) -> CommandResult:87 """Execute simulated S3 commands."""88 bucket_name = args.get("bucket") or self._extract_bucket_name(command)89 90 if "get-public-access-block" in command:91 if bucket_name in self.cloud_state["s3_buckets"]:92 resource = self.cloud_state["s3_buckets"][bucket_name]93 if resource.public:94 output = f'{{"PublicAccessBlockConfiguration": {{"BlockPublicAcls": false}}}}'95 else:96 output = (97 '{"PublicAccessBlockConfiguration": {"BlockPublicAcls": true}}'98 )99 return CommandResult(success=True, output=output, command=command)100 return CommandResult(101 success=False,102 output="",103 error=f"NoSuchBucket: The specified bucket does not exist",104 command=command,105 )106 107 elif "put-public-access-block" in command:108 if bucket_name in self.cloud_state["s3_buckets"]:109 self.cloud_state["s3_buckets"][bucket_name].public = False110 self.cloud_state["s3_buckets"][bucket_name].properties[111 "public_access_blocked"112 ] = True113 self._log_execution(114 command, {"bucket": bucket_name, "action": "block_public_access"}115 )116 return CommandResult(117 success=True,118 output=f"Public access blocked for bucket: {bucket_name}",119 command=command,120 )121 return CommandResult(122 success=False,123 output="",124 error=f"NoSuchBucket: The specified bucket does not exist",125 command=command,126 )127 128 elif "list-buckets" in command:129 buckets = list(self.cloud_state["s3_buckets"].keys())130 output = '{"Buckets": ' + str([{"Name": b} for b in buckets]) + "}"131 return CommandResult(success=True, output=output, command=command)132 133 return CommandResult(134 success=False,135 output="",136 error=f"Unknown S3 command: {command}",137 command=command,138 )139 140 def _execute_iam_command(self, command: str, args: Dict[str, Any]) -> CommandResult:141 """Execute simulated IAM commands."""142 user_name = (143 args.get("user-name")144 or args.get("UserName")145 or self._extract_username(command)146 )147 148 if "get-user" in command:149 if user_name in self.cloud_state["iam_users"]:150 resource = self.cloud_state["iam_users"][user_name]151 status = "Disabled" if resource.disabled else "Active"152 output = (153 f'{{"User": {{"UserName": "{user_name}", "Status": "{status}"}}}}'154 )155 return CommandResult(success=True, output=output, command=command)156 return CommandResult(157 success=False,158 output="",159 error=f"NoSuchEntity: The user {user_name} does not exist",160 command=command,161 )162 163 elif "update-user" in command or "update-user" in command:164 status = args.get("status") or "Active"165 if user_name in self.cloud_state["iam_users"]:166 self.cloud_state["iam_users"][user_name].disabled = status == "Disabled"167 self._log_execution(command, {"user": user_name, "status": status})168 return CommandResult(169 success=True,170 output=f"User {user_name} updated to status: {status}",171 command=command,172 )173 return CommandResult(174 success=False,175 output="",176 error=f"NoSuchEntity: The user {user_name} does not exist",177 command=command,178 )179 180 elif "list-users" in command:181 users = list(self.cloud_state["iam_users"].keys())182 output = '{"Users": ' + str([{"UserName": u} for u in users]) + "}"183 return CommandResult(success=True, output=output, command=command)184 185 return CommandResult(186 success=False,187 output="",188 error=f"Unknown IAM command: {command}",189 command=command,190 )191 192 def _execute_ec2_command(self, command: str, args: Dict[str, Any]) -> CommandResult:193 """Execute simulated EC2 commands."""194 if "describe-instances" in command:195 instances = []196 for inst_id, inst in self.cloud_state["ec2_instances"].items():197 instances.append(198 {199 "InstanceId": inst_id,200 "State": {201 "Name": "running" if not inst.disabled else "stopped"202 },203 **inst.properties,204 }205 )206 output = '{"Reservations": [{"Instances": ' + str(instances) + "}]}"207 return CommandResult(success=True, output=output, command=command)208 209 return CommandResult(210 success=False,211 output="",212 error=f"Unknown EC2 command: {command}",213 command=command,214 )215 216 def _execute_security_group_command(217 self, command: str, args: Dict[str, Any]218 ) -> CommandResult:219 """Execute simulated security group commands."""220 group_id = args.get("GroupId") or "sg-default"221 222 if group_id in self.cloud_state["security_groups"]:223 sg = self.cloud_state["security_groups"][group_id]224 output = str({"SecurityGroups": [{"GroupId": group_id, **sg.properties}]})225 return CommandResult(success=True, output=output, command=command)226 227 return CommandResult(228 success=False,229 output="",230 error=f"InvalidGroup.NotFound: Security group {group_id} not found",231 command=command,232 )233 234 def _extract_bucket_name(self, command: str) -> Optional[str]:235 """Extract bucket name from command string."""236 match = re.search(r"--bucket\s+(\S+)", command)237 if match:238 return match.group(1)239 match = re.search(r"bucket[/\s]+(\S+)", command, re.IGNORECASE)240 if match:241 return match.group(1)242 return None243 244 def _extract_username(self, command: str) -> Optional[str]:245 """Extract username from command string."""246 match = re.search(r"--user-name\s+(\S+)", command)247 if match:248 return match.group(1)249 match = re.search(r"user[/\s]+(\S+)", command, re.IGNORECASE)250 if match:251 return match.group(1)252 return None253 254 def _log_execution(self, command: str, metadata: Dict[str, Any]):255 """Log command execution."""256 self.execution_log.append(257 {258 "command": command,259 "metadata": metadata,260 "timestamp": datetime.now().isoformat(),261 "execution_id": self._execution_count,262 }263 )264 265 def get_state(266 self, resource_type: str, resource_id: str267 ) -> Optional[CloudResource]:268 """Get current state of a resource."""269 if resource_type in self.cloud_state:270 return self.cloud_state[resource_type].get(resource_id)271 return None272 273 def update_state(274 self, resource_type: str, resource_id: str, properties: Dict[str, Any]275 ):276 """Update simulated resource state."""277 if resource_type not in self.cloud_state:278 self.cloud_state[resource_type] = {}279 280 if resource_id in self.cloud_state[resource_type]:281 self.cloud_state[resource_type][resource_id].properties.update(properties)282 else:283 self.cloud_state[resource_type][resource_id] = CloudResource(284 resource_type=resource_type,285 resource_id=resource_id,286 properties=properties,287 )288 289 def add_bucket(self, bucket_name: str, public: bool = False):290 """Add a simulated S3 bucket to state."""291 self.cloud_state["s3_buckets"][bucket_name] = CloudResource(292 resource_type="s3",293 resource_id=bucket_name,294 properties={"Name": bucket_name},295 public=public,296 )297 298 def add_user(self, user_name: str, disabled: bool = False):299 """Add a simulated IAM user to state."""300 self.cloud_state["iam_users"][user_name] = CloudResource(301 resource_type="iam",302 resource_id=user_name,303 properties={"UserName": user_name},304 disabled=disabled,305 )306 307 def add_ec2_instance(self, instance_id: str, properties: Dict[str, Any] = None):308 """Add a simulated EC2 instance to state."""309 self.cloud_state["ec2_instances"][instance_id] = CloudResource(310 resource_type="ec2",311 resource_id=instance_id,312 properties=properties or {},313 )314 315 def add_security_group(self, group_id: str, properties: Dict[str, Any] = None):316 """Add a simulated security group to state."""317 self.cloud_state["security_groups"][group_id] = CloudResource(318 resource_type="security_group",319 resource_id=group_id,320 properties=properties or {},321 )322 323 def simulate_delay(self, min_seconds: float = 0.1, max_seconds: float = 0.5):324 """Add realistic delay to command execution."""325 import random326 327 delay = random.uniform(min_seconds, max_seconds)328 time.sleep(delay)329 330 def generate_audit_log(self) -> List[Dict[str, Any]]:331 """Return log of all executed commands."""332 return self.execution_log.copy()333 334 def reset(self):335 """Reset simulator state."""336 self.execution_log = []337 self.cloud_state = {338 "s3_buckets": {},339 "iam_users": {},340 "ec2_instances": {},341 "security_groups": {},342 }343 self._execution_count = 0344 345 def get_execution_summary(self) -> Dict[str, Any]:346 """Get summary of command executions."""347 return {348 "total_executions": self._execution_count,349 "execution_log_size": len(self.execution_log),350 "resources": {351 resource_type: len(resources)352 for resource_type, resources in self.cloud_state.items()353 },354 }355