main.py

"""CloudGuard-S3-Auditor — OpenEnv Environment (v0.1.0)

FastAPI server exposing /reset, /step, /grade endpoints on port 7860.
"""

from __future__ import annotations

import random
import uuid
from contextlib import asynccontextmanager
from typing import Any

from fastapi import FastAPI, HTTPException
from pydantic import BaseModel, Field

from models import (
    Action,
    ActionType,
    BucketObservation,
    BucketPolicy,
    EncryptionType,
    GradeResult,
    Observation,
    S3Bucket,
    StepResult,
)

# ── Task Definitions ─────────────────────────────────────────────────────────

TASK_REGISTRY: dict[str, dict] = {
    "critical_leak": {
        "difficulty": "easy",
        "max_steps": 3,
        "generator": "_gen_critical_leak",
    },
    "audit_and_protect": {
        "difficulty": "medium",
        "max_steps": 15,
        "generator": "_gen_audit_and_protect",
    },
    "compliance_sweep": {
        "difficulty": "hard",
        "max_steps": 30,
        "generator": "_gen_compliance_sweep",
    },
}

# ── Bucket Generators ────────────────────────────────────────────────────────

_PII_NAMES = [
    "user-records", "customer-pii", "employee-ssn", "payment-cards",
    "health-data", "tax-forms", "passport-scans", "credit-reports",
    "insurance-claims", "hr-documents",
]
_PUBLIC_ASSET_NAMES = [
    "website-css", "website-images", "cdn-assets", "static-js",
    "marketing-banners", "favicon-assets",
]
_NORMAL_NAMES = [
    "app-logs", "build-artifacts", "internal-docs", "ml-training-data",
    "analytics-raw", "backup-2024", "staging-dump", "config-snapshots",
    "ci-cache", "vendor-contracts",
]


def _gen_critical_leak() -> list[S3Bucket]:
    return [
        S3Bucket(
            name="customer-pii",
            policy=BucketPolicy.PUBLIC,
            encryption=EncryptionType.NONE,
            contains_pii=True,
            is_required_public=False,
        )
    ]


def _gen_audit_and_protect() -> list[S3Bucket]:
    buckets: list[S3Bucket] = []

    # 3 PII buckets — public (need fixing)
    for name in random.sample(_PII_NAMES, 3):
        buckets.append(
            S3Bucket(
                name=name,
                policy=BucketPolicy.PUBLIC,
                encryption=EncryptionType.NONE,
                contains_pii=True,
                is_required_public=False,
            )
        )

    # 2 required-public buckets — must stay public
    for name in random.sample(_PUBLIC_ASSET_NAMES, 2):
        buckets.append(
            S3Bucket(
                name=name,
                policy=BucketPolicy.PUBLIC,
                encryption=EncryptionType.NONE,
                contains_pii=False,
                is_required_public=True,
            )
        )

    # 5 normal private buckets (filler, already fine)
    for name in random.sample(_NORMAL_NAMES, 5):
        buckets.append(
            S3Bucket(
                name=name,
                policy=BucketPolicy.PRIVATE,
                encryption=EncryptionType.AES256,
                contains_pii=False,
                is_required_public=False,
            )
        )

    random.shuffle(buckets)
    return buckets


def _gen_compliance_sweep() -> list[S3Bucket]:
    buckets: list[S3Bucket] = []

    # 5 PII buckets — public, unencrypted (need: private + encrypted)
    for name in random.sample(_PII_NAMES, 5):
        buckets.append(
            S3Bucket(
                name=name,
                policy=BucketPolicy.PUBLIC,
                encryption=EncryptionType.NONE,
                contains_pii=True,
                is_required_public=False,
            )
        )

    # 4 required-public (traps — must stay public)
    for name in random.sample(_PUBLIC_ASSET_NAMES, 4):
        buckets.append(
            S3Bucket(
                name=name,
                policy=BucketPolicy.PUBLIC,
                encryption=EncryptionType.NONE,
                contains_pii=False,
                is_required_public=True,
            )
        )

    # 6 normal buckets — private but unencrypted (need encryption only)
    for name in random.sample(_NORMAL_NAMES, 6):
        buckets.append(
            S3Bucket(
                name=name,
                policy=BucketPolicy.PRIVATE,
                encryption=EncryptionType.NONE,
                contains_pii=False,
                is_required_public=False,
            )
        )

    # 5 normal buckets — already compliant (filler)
    remaining = [n for n in _NORMAL_NAMES if n not in [b.name for b in buckets]]
    for name in random.sample(remaining, min(5, len(remaining))):
        buckets.append(
            S3Bucket(
                name=name,
                policy=BucketPolicy.PRIVATE,
                encryption=EncryptionType.AES256,
                contains_pii=False,
                is_required_public=False,
            )
        )

    random.shuffle(buckets)
    return buckets


_GENERATORS = {
    "critical_leak": _gen_critical_leak,
    "audit_and_protect": _gen_audit_and_protect,
    "compliance_sweep": _gen_compliance_sweep,
}


# ── Environment ──────────────────────────────────────────────────────────────

class S3AuditorEnv:
    """Core environment logic."""

    def __init__(self) -> None:
        self.buckets: list[S3Bucket] = []
        self.task_id: str = ""
        self.max_steps: int = 0
        self.current_step: int = 0
        self.done: bool = False
        self.rewards: list[float] = []
        self._initial_issues: int = 0

    # ── reset ────────────────────────────────────────────────────────────

    def reset(self, task_id: str, seed: int | None = None) -> Observation:
        if task_id not in TASK_REGISTRY:
            raise ValueError(f"Unknown task: {task_id}")

        if seed is not None:
            random.seed(seed)

        meta = TASK_REGISTRY[task_id]
        self.task_id = task_id
        self.max_steps = meta["max_steps"]
        self.current_step = 0
        self.done = False
        self.rewards = []
        self.buckets = _GENERATORS[task_id]()
        self._initial_issues = self._count_issues()

        return self._observe("Environment reset. Audit the buckets.")

    # ── step ─────────────────────────────────────────────────────────────

    def step(self, action: Action) -> StepResult:
        if self.done:
            return StepResult(
                observation=self._observe("Episode already done."),
                reward=0.0,
                done=True,
            )

        self.current_step += 1
        reward = 0.0
        msg = ""

        if action.action_type == ActionType.NOOP:
            msg = "No operation."

        else:
            bucket = self._find_bucket(action.bucket_name)
            if bucket is None:
                self.rewards.append(0.0)
                return StepResult(
                    observation=self._observe(f"Bucket '{action.bucket_name}' not found."),
                    reward=0.0,
                    done=False,
                    info={"error": "bucket_not_found"},
                )

            if action.action_type == ActionType.MAKE_PRIVATE:
                if bucket.is_required_public:
                    reward = -1.0
                    msg = f"CRITICAL: '{bucket.name}' is a required-public asset! Production broken."
                elif bucket.policy == BucketPolicy.PUBLIC and bucket.contains_pii:
                    bucket.policy = BucketPolicy.PRIVATE
                    reward = 0.2
                    msg = f"Secured PII bucket '{bucket.name}'."
                elif bucket.policy == BucketPolicy.PUBLIC:
                    bucket.policy = BucketPolicy.PRIVATE
                    reward = 0.0
                    msg = f"Made '{bucket.name}' private (no PII — neutral)."
                else:
                    msg = f"'{bucket.name}' is already private."

            elif action.action_type == ActionType.MAKE_PUBLIC:
                if bucket.contains_pii and bucket.policy == BucketPolicy.PRIVATE:
                    reward = -1.0
                    msg = f"CRITICAL: Exposed PII bucket '{bucket.name}' to public!"
                else:
                    bucket.policy = BucketPolicy.PUBLIC
                    msg = f"Made '{bucket.name}' public."

            elif action.action_type == ActionType.ENABLE_ENCRYPTION:
                algo = action.encryption_algo or EncryptionType.AES256
                if bucket.encryption == EncryptionType.NONE:
                    bucket.encryption = algo
                    reward = 0.2
                    msg = f"Enabled {algo.value} encryption on '{bucket.name}'."
                else:
                    msg = f"'{bucket.name}' already encrypted ({bucket.encryption.value})."

        self.rewards.append(reward)

        # Check terminal conditions
        if self.current_step >= self.max_steps:
            self.done = True
            msg += " Max steps reached."
        elif self._count_issues() == 0:
            self.done = True
            # Minimum-steps bonus
            min_actions = self._initial_issues
            if self.current_step <= min_actions and all(r >= 0 for r in self.rewards):
                bonus = 0.4
                self.rewards[-1] += bonus
                reward += bonus
                msg += " BONUS: Completed in minimum steps!"
            msg += " All issues resolved."

        return StepResult(
            observation=self._observe(msg),
            reward=round(reward, 4),
            done=self.done,
        )

    # ── grading ──────────────────────────────────────────────────────────

    def grade(self) -> GradeResult:
        """Evaluate final state. Returns normalized score in [0, 1]."""
        issues = self._count_issues()
        broken = self._count_broken_production()

        # Max possible positive reward:
        #   initial_issues * 0.2 + 0.4 bonus = theoretical max
        max_raw = self._initial_issues * 0.2 + 0.4
        raw_sum = sum(self.rewards)

        # Normalize to [0, 1]
        if max_raw > 0:
            score = max(0.0, min(1.0, raw_sum / max_raw))
        else:
            score = 1.0 if issues == 0 else 0.0

        success = issues == 0 and broken == 0

        return GradeResult(
            success=success,
            score=round(score, 4),
            details={
                "remaining_issues": issues,
                "broken_production": broken,
                "total_steps": self.current_step,
                "raw_reward_sum": round(raw_sum, 4),
                "rewards": [round(r, 4) for r in self.rewards],
            },
        )

    # ── helpers ──────────────────────────────────────────────────────────

    def _find_bucket(self, name: str | None) -> S3Bucket | None:
        if name is None:
            return None
        for b in self.buckets:
            if b.name == name:
                return b
        return None

    def _count_issues(self) -> int:
        """Count buckets that still need fixing."""
        count = 0
        for b in self.buckets:
            if b.contains_pii and b.policy == BucketPolicy.PUBLIC:
                count += 1
            if b.encryption == EncryptionType.NONE and not b.is_required_public:
                count += 1
        return count

    def _count_broken_production(self) -> int:
        return sum(
            1 for b in self.buckets
            if b.is_required_public and b.policy == BucketPolicy.PRIVATE
        )

    def _observe(self, message: str = "") -> Observation:
        return Observation(
            buckets=[
                BucketObservation(
                    name=b.name,
                    policy=b.policy,
                    encryption=b.encryption,
                    contains_pii=b.contains_pii,
                    is_required_public=b.is_required_public,
                )
                for b in self.buckets
            ],
            step_number=self.current_step,
            max_steps=self.max_steps,
            done=self.done,
            message=message,
        )


# ── FastAPI Server ───────────────────────────────────────────────────────────

env = S3AuditorEnv()


class ResetRequest(BaseModel):
    task_id: str = "critical_leak"
    seed: int | None = None


class StepRequest(BaseModel):
    action: str = Field(..., description="Action string, e.g. 'make_private:customer-pii'")


@asynccontextmanager
async def lifespan(app: FastAPI):
    yield

app = FastAPI(
    title="CloudGuard-S3-Auditor",
    version="0.1.0",
    lifespan=lifespan,
)


@app.get("/health")
async def health():
    return {"status": "ok", "env": "CloudGuard-S3-Auditor", "version": "0.1.0"}


@app.post("/reset", response_model=Observation)
async def reset(req: ResetRequest):
    try:
        obs = env.reset(task_id=req.task_id, seed=req.seed)
        return obs
    except ValueError as e:
        raise HTTPException(status_code=400, detail=str(e))


@app.post("/step", response_model=StepResult)
async def step(req: StepRequest):
    try:
        action = Action.from_str(req.action)
    except (ValueError, KeyError) as e:
        raise HTTPException(status_code=400, detail=f"Invalid action: {e}")
    result = env.step(action)
    return result


@app.post("/grade", response_model=GradeResult)
async def grade():
    return env.grade()


@app.get("/tasks")
async def list_tasks():
    return {
        tid: {"difficulty": meta["difficulty"], "max_steps": meta["max_steps"]}
        for tid, meta in TASK_REGISTRY.items()
    }


if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=7860)