Spaces:
Sleeping
Sleeping
File size: 2,351 Bytes
b7bab2b 2de3485 0955106 2de3485 0955106 2de3485 0bbdc65 2de3485 0955106 2de3485 b7bab2b 2de3485 7d002aa b7bab2b 0955106 b7bab2b 0955106 b7bab2b 0955106 b7bab2b 2de3485 941ab5e 0955106 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 | from typing import List, Optional, Union, Dict
from pydantic import BaseModel, Field
# --- Observation Models ---
class SecurityGroupRule(BaseModel):
port: int
cidr: str
class SecurityGroup(BaseModel):
id: str
name: str
ingress_rules: List[SecurityGroupRule]
class S3Bucket(BaseModel):
name: str
encrypted: bool
class RDSInstance(BaseModel):
id: str
engine: str
encrypted: bool
class EBSVolume(BaseModel):
id: str
encrypted: bool
class IAMPolicy(BaseModel):
id: str
name: str
document: str
class CloudObservation(BaseModel):
security_groups: List[SecurityGroup]
s3_buckets: List[S3Bucket]
rds_instances: List[RDSInstance] = []
ebs_volumes: List[EBSVolume] = []
iam_policies: List[IAMPolicy]
task_description: str = "Perform a cloud security audit and remediate vulnerabilities."
vulnerability_manifest: Dict[str, int] = {} # e.g. {"sg_vulns": 4, "s3_vulns": 3}
message: str = "Cloud resources loaded."
reward: float = 0.0
health_score: float = 1.0 # 0.0 to 1.0 (AVAILABILITY)
done: bool = False
info: dict = {}
# --- Unified Action Model ---
class CloudAction(BaseModel):
"""
Unified action model to resolve union validation issues.
All specialized fields are optional.
"""
action_type: str # "audit", "fix_sg", "remediate_all_in_sg", "enable_s3_enc", "enable_rds_enc", "enable_ebs_enc", "update_iam", "submit"
# Target identifiers
sg_id: Optional[str] = None
bucket_name: Optional[str] = None
policy_id: Optional[str] = None
rds_id: Optional[str] = None
ebs_id: Optional[str] = None
# fix_sg specific fields
port: Optional[int] = None
cidr_to_remove: Optional[str] = None
# update_iam fields
new_document: Optional[str] = None
# submit fields
findings: Optional[List[str]] = None
# --- State Model ---
class CloudState(BaseModel):
task_name: str
step_count: int
max_steps: int
remediated_count: int
security_groups: List[SecurityGroup] = []
s3_buckets: List[S3Bucket] = []
rds_instances: List[RDSInstance] = []
ebs_volumes: List[EBSVolume] = []
iam_policies: List[IAMPolicy] = []
vulnerability_manifest: Dict[str, int] = {}
required_iam_perms: Dict[str, str] = {}
health_score: float = 1.0
|