models.py

"""Pydantic V2 models for CloudGuard-S3-Auditor environment."""

from __future__ import annotations

import enum
from typing import Literal

from pydantic import BaseModel, Field, field_validator


# ── Enums ────────────────────────────────────────────────────────────────────

class BucketPolicy(str, enum.Enum):
    PUBLIC = "public"
    PRIVATE = "private"


class EncryptionType(str, enum.Enum):
    NONE = "none"
    AES256 = "aes256"
    AWS_KMS = "aws-kms"


class ActionType(str, enum.Enum):
    MAKE_PRIVATE = "make_private"
    MAKE_PUBLIC = "make_public"
    ENABLE_ENCRYPTION = "enable_encryption"
    NOOP = "noop"


# ── Bucket State ─────────────────────────────────────────────────────────────

class S3Bucket(BaseModel):
    """Single S3 bucket in the simulated environment."""

    name: str = Field(..., description="Bucket identifier, e.g. 'website-css' or 'user-pii-db'")
    policy: BucketPolicy = Field(default=BucketPolicy.PUBLIC)
    encryption: EncryptionType = Field(default=EncryptionType.NONE)
    contains_pii: bool = Field(default=False)
    is_required_public: bool = Field(default=False)

    def policy_json(self) -> dict:
        return {
            "Version": "2012-10-17",
            "Statement": [
                {
                    "Effect": "Allow" if self.policy == BucketPolicy.PUBLIC else "Deny",
                    "Principal": "*",
                    "Action": "s3:GetObject",
                    "Resource": f"arn:aws:s3:::{self.name}/*",
                }
            ],
        }


# ── Action ───────────────────────────────────────────────────────────────────

class Action(BaseModel):
    """Agent action submitted each step."""

    action_type: ActionType
    bucket_name: str | None = Field(
        default=None,
        description="Target bucket name (required for all actions except noop)",
    )
    encryption_algo: EncryptionType | None = Field(
        default=None,
        description="Encryption algorithm (required only for enable_encryption)",
    )

    @field_validator("bucket_name")
    @classmethod
    def require_bucket_for_non_noop(cls, v: str | None, info) -> str | None:
        if info.data.get("action_type") != ActionType.NOOP and v is None:
            raise ValueError("bucket_name is required for non-noop actions")
        return v

    def to_str(self) -> str:
        if self.action_type == ActionType.NOOP:
            return "noop"
        if self.action_type == ActionType.ENABLE_ENCRYPTION:
            algo = self.encryption_algo.value if self.encryption_algo else "aes256"
            return f"enable_encryption:{self.bucket_name}:{algo}"
        return f"{self.action_type.value}:{self.bucket_name}"

    @classmethod
    def from_str(cls, raw: str) -> "Action":
        raw = raw.strip()
        if raw == "noop":
            return cls(action_type=ActionType.NOOP)

        parts = raw.split(":")
        action_type = ActionType(parts[0])

        if action_type == ActionType.ENABLE_ENCRYPTION:
            if len(parts) < 3:
                raise ValueError("enable_encryption requires bucket_name and algo")
            return cls(
                action_type=action_type,
                bucket_name=parts[1],
                encryption_algo=EncryptionType(parts[2]),
            )

        return cls(action_type=action_type, bucket_name=parts[1])


# ── Observation ──────────────────────────────────────────────────────────────

class BucketObservation(BaseModel):
    """Visible state of a single bucket (sent to agent)."""

    name: str
    policy: BucketPolicy
    encryption: EncryptionType
    contains_pii: bool
    is_required_public: bool


class Observation(BaseModel):
    """Full observation returned to agent each step."""

    buckets: list[BucketObservation]
    step_number: int = Field(ge=0)
    max_steps: int = Field(ge=1)
    done: bool = False
    message: str = ""


# ── Reward / Step Result ─────────────────────────────────────────────────────

class StepResult(BaseModel):
    """Result of a single env.step() call."""

    observation: Observation
    reward: float = Field(ge=-10.0, le=10.0)
    done: bool = False
    info: dict = Field(default_factory=dict)


# ── Grading ──────────────────────────────────────────────────────────────────

class GradeResult(BaseModel):
    """Output of the SuccessGrader."""

    success: bool
    score: float = Field(ge=0.0, le=1.0)
    details: dict = Field(default_factory=dict)