File size: 2,048 Bytes
f1a1961 547b872 f1a1961 | 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 | import requests
import json
BASE_URL = "http://localhost:7860"
def run_baseline_audit(task_id="easy"):
print(f"--- Running Baseline for Task: {task_id} ---")
# 1. Reset environment
response = requests.post(f"{BASE_URL}/reset", json={"task_id": task_id})
if response.status_code != 200:
print(f"Failed to reset: {response.text}")
return
obs_data = response.json()
obs = obs_data.get("observation", {})
print(f"Observation Info: {obs.get('info')}")
# 2. List S3 buckets
# Note: wrapping in "action" key to avoid collision with 'action' field in CloudAction
action_payload = {
"action": {
"action": "list",
"resource_type": "s3"
}
}
response = requests.post(f"{BASE_URL}/step", json=action_payload)
if response.status_code != 200:
print(f"Failed on step: {response.text}")
return
step_result = response.json()
obs = step_result.get("observation", {})
resources = obs.get("resources", [])
print(f"Discovered {len(resources)} S3 buckets.")
# 3. Logic to identify public prod buckets
public_prod_buckets = []
for r in resources:
if r.get("public") and r.get("tags", {}).get("env") == "prod":
public_prod_buckets.append(r["id"])
print(f"Identified Public Prod Buckets: {public_prod_buckets}")
# 4. Submit answer
submit_payload = {
"action": {
"action": "submit",
"answer": ",".join(public_prod_buckets)
}
}
response = requests.post(f"{BASE_URL}/step", json=submit_payload)
step_result = response.json()
obs = step_result.get("observation", {})
reward = step_result.get("reward", 0.0)
done = step_result.get("done", False)
print(f"Final Reward: {reward}")
print(f"Done: {done}")
print(f"Info: {obs.get('info')}")
if __name__ == "__main__":
try:
run_baseline_audit("easy")
except Exception as e:
print(f"Error: {e}")
|