inference.py

import sys
from typing import Any
import requests

# ==========================================
# CONFIGURATION
# ==========================================
MY_SPACE_URL = "https://shiggii-s3-security-agent-env.hf.space" 
ENV_NAME = "CloudGuard-S3-Auditor"
MODEL_NAME = "rule-based-baseline"

# ==========================================
# LOGGING UTILITIES
# ==========================================
def log_start(task: str) -> None:
    print(f"[START] task={task} env={ENV_NAME} model={MODEL_NAME}", flush=True)

def log_step(step: int, action: str, reward: float, done: bool, error: str | None) -> None:
    err_str = error if error else "null"
    print(
        f"[STEP]  step={step} action={action} reward={reward:.2f} "
        f"done={str(done).lower()} error={err_str}",
        flush=True,
    )

def log_end(success: bool, steps: int, score: float, rewards: list[float]) -> None:
    rewards_str = ",".join(f"{r:.2f}" for r in rewards)
    print(
        f"[END]   success={str(success).lower()} steps={steps} "
        f"score={score:.4f} rewards={rewards_str}",
        flush=True,
    )

# ==========================================
# REFINED BASELINE POLICY
# ==========================================
def decide_actions(obs: dict) -> list[str]:
    actions = []
    if "buckets" not in obs: return ["noop"]

    for bucket in obs["buckets"]:
        name = bucket.get("name")
        is_public = bucket.get("policy") == "public"
        has_pii = bucket.get("contains_pii") is True
        is_req_public = bucket.get("is_required_public") is True
        is_encrypted = bucket.get("encryption") != "none"

        # 1. Priority: Secure PII (Confirmed working with make_private)
        if is_public and has_pii:
            actions.append(f"make_private:{name}")

        # 2. Priority: Compliance Encryption (Adding algorithm parameter to fix Error 400)
        if not is_encrypted and not is_req_public:
            actions.append(f"enable_encryption:{name}:aes256")

    return actions if actions else ["noop"]

# ==========================================
# EXECUTION ENGINE
# ==========================================
def run_task(task_id: str, base_url: str):
    log_start(task_id)
    base_url = base_url.rstrip('/')

    # 1. Reset
    try:
        resp = requests.post(f"{base_url}/reset", json={"task_id": task_id}, timeout=15)
        resp.raise_for_status()
        obs = resp.json()
    except Exception as e:
        print(f"❌ Connection Error: {e}")
        return

    all_rewards = []
    step_num = 0
    plan = decide_actions(obs)

    # 2. Execute Actions
    for action_str in plan:
        step_num += 1
        try:
            resp = requests.post(f"{base_url}/step", json={"action": action_str}, timeout=10)
            
            # Fallback: If colon fails, try underscore, then Uppercase
            if resp.status_code == 400:
                alt_action = action_str.replace(":", "_")
                resp = requests.post(f"{base_url}/step", json={"action": alt_action}, timeout=10)
                if resp.status_code == 200:
                    action_str = alt_action

            if resp.status_code != 200:
                log_step(step_num, action_str, 0.0, False, f"Err {resp.status_code}")
                all_rewards.append(0.0)
                continue

            result = resp.json()
            reward = result.get("reward", 0.0)
            done = result.get("done", False)
            
            all_rewards.append(reward)
            log_step(step_num, action_str, reward, done, None)
            if done: break
        except Exception as e:
            log_step(step_num, action_str, 0.0, False, str(e))

    # 3. Final Grade
    try:
        grade_resp = requests.post(f"{base_url}/grade", timeout=10)
        grade = grade_resp.json()
        log_end(
            success=grade.get("success", False), 
            steps=step_num, 
            score=grade.get("score", 0.0), 
            rewards=all_rewards
        )
    except:
        print("❌ Final Grade failed.")

# ==========================================
# RUN
# ==========================================
print("🚀 Running Final CloudGuard Baseline...\n")
for t in ["critical_leak", "audit_and_protect", "compliance_sweep"]:
    run_task(t, MY_SPACE_URL)
    print("-" * 50)