import asyncio import random from typing import Optional, List, Dict from types import SimpleNamespace from fastapi import FastAPI from openenv.core.env_server import Environment from models import MyEnvV4Observation, MyEnvV4Action, URLInfo class MyEnvV4Env(Environment): def __init__(self): super().__init__() self.dataset = self._generate_sophisticated_dataset() # For evaluation reproducibility, we could shuffle, but for benchmark stability, we keep order self.current_step = 0 def _generate_sophisticated_dataset(self): """ Expanded dataset with 15 samples across 3 difficulty levels. Covers Clean, Spam, and Phishing (Digital Seduction). """ base_data = [ # LEVEL 1: CLEAR CASES (5) { "sender": "registrar@manipal.edu", "subject": "Semester Registration Open", "body": "Dear student, please register for the Fall semester via the student portal.", "raw_headers": "Received: from mail.manipal.edu (14.139.161.20) by mx.google.com; SPF: pass; DKIM: pass;", "auth_results": {"SPF": "pass", "DKIM": "pass", "DMARC": "pass"}, "urls": [], "label": "INBOX", "difficulty": 1 }, { "sender": "no-reply@amazon.com", "subject": "Your Order #123-4567 has shipped", "body": "Track your package delivery status in your Amazon account.", "raw_headers": "Received: from a9-12.smtp-out.amazonses.com... SPF: pass; DKIM: pass;", "auth_results": {"SPF": "pass", "DKIM": "pass", "DMARC": "pass"}, "urls": [{"url": "https://amazon.com/track", "is_shortened": False, "domain_age_days": 8000, "has_ssl": True, "reputation_score": 1.0}], "label": "INBOX", "difficulty": 1 }, { "sender": "win-money@lottery-global.net", "subject": "YOU WON!! $1,000,000", "body": "Congratulations! You have been selected as our winner. CLAIM YOUR $1M NOW!", "raw_headers": "Received: from unknown-relay.co (103.22.1.5)... SPF: none; DKIM: fail;", "auth_results": {"SPF": "none", "DKIM": "fail", "DMARC": "none"}, "urls": [{"url": "http://get-cash-free.net/claim", "is_shortened": False, "domain_age_days": 2, "has_ssl": False, "reputation_score": 0.1}], "label": "SPAM", "difficulty": 1 }, { "sender": "viagra-sales@pharma-dealz.biz", "subject": "Cheapest Meds Online", "body": "Buy now and save 90% on all prescription drugs. No prescription needed!", "raw_headers": "Received: from botnet-node.ru... SPF: softfail;", "auth_results": {"SPF": "softfail", "DKIM": "none", "DMARC": "none"}, "urls": [{"url": "http://cheap-rx.biz", "is_shortened": False, "domain_age_days": 15, "has_ssl": False, "reputation_score": 0.05}], "label": "SPAM", "difficulty": 1 }, { "sender": "support@netflix-security.com", "subject": "Update Payment Method", "body": "Your Netflix subscription has expired. Click here to login and update billing.", "raw_headers": "Received: from suspicious-vps.com... SPF: fail; DMARC: fail;", "auth_results": {"SPF": "fail", "DKIM": "none", "DMARC": "fail"}, "urls": [{"url": "https://bit.ly/fake-netflix-login", "is_shortened": True, "domain_age_days": 3, "has_ssl": True, "reputation_score": 0.02}], "label": "QUARANTINE", "difficulty": 1 }, # LEVEL 2: NUANCED / MARKETING / LEGIT BUT ANNOYING (5) { "sender": "news@internshala-mail.com", "subject": "New Internships in Manipal", "body": "Check out these new opportunities for CSE students. Apply today!", "raw_headers": "Received: from mktg.server.com... SPF: pass; DKIM: pass;", "auth_results": {"SPF": "pass", "DKIM": "pass", "DMARC": "pass"}, "urls": [{"url": "https://internshala.com/n/123", "is_shortened": False, "domain_age_days": 2500, "has_ssl": True, "reputation_score": 0.95}], "label": "SPAM", # Triaged as Promotions/Spam "difficulty": 2 }, { "sender": "marketing@swiggy.in", "subject": "50% OFF your next meal!", "body": "Hungry? Use code HUNGRY50 at checkout.", "raw_headers": "Received: from swiggy-mail.in... SPF: pass; DKIM: pass;", "auth_results": {"SPF": "pass", "DKIM": "pass", "DMARC": "pass"}, "urls": [], "label": "SPAM", "difficulty": 2 }, { "sender": "hr@startup-hiring.co", "subject": "Interview Invitation", "body": "We saw your profile on LinkedIn and want to chat about a role.", "raw_headers": "Received: from linkedin-referral.com... SPF: neutral;", "auth_results": {"SPF": "neutral", "DKIM": "none", "DMARC": "none"}, "urls": [{"url": "https://startup-hiring.co/apply", "is_shortened": False, "domain_age_days": 45, "has_ssl": True, "reputation_score": 0.6}], "label": "INBOX", # Legitimate cold outreach "difficulty": 2 }, { "sender": "alert@banking-secure.net", "subject": "Suspicious Activity Detected", "body": "We detected an unusual login to your account from Russia. Please verify.", "raw_headers": "Received: from spoofed-host.com... SPF: softfail; DMARC: none;", "auth_results": {"SPF": "softfail", "DKIM": "none", "DMARC": "none"}, "urls": [{"url": "https://t.co/secure-bank-login", "is_shortened": True, "domain_age_days": 10, "has_ssl": True, "reputation_score": 0.3}], "label": "QUARANTINE", "difficulty": 2 }, { "sender": "noreply@github.com", "subject": "[GitHub] A personal access token has been added", "body": "A new personal access token was added to your account. If this wasn't you, click here.", "raw_headers": "Received: from out-21.smtp.github.com... SPF: pass; DKIM: pass;", "auth_results": {"SPF": "pass", "DKIM": "pass", "DMARC": "pass"}, "urls": [{"url": "https://github.com/settings/tokens", "is_shortened": False, "domain_age_days": 6000, "has_ssl": True, "reputation_score": 1.0}], "label": "INBOX", "difficulty": 2 }, # LEVEL 3: DIGITAL SEDUCTION (PHISHING / SPEAR PHISHING) (5) { "sender": "dean.office@manipal-edu.in", # Typo-squatted (.in instead of .edu) "subject": "Urgent: Faculty Grievance Report", "body": "A report has been filed against your department. Review the grievances here immediately.", "raw_headers": "Received: from rogue-vps.xyz (45.33.22.11) by mx.google.com; SPF: fail; DMARC: fail;", "auth_results": {"SPF": "fail", "DKIM": "none", "DMARC": "fail"}, "urls": [{"url": "https://bit.ly/3xYz1-grievance", "is_shortened": True, "domain_age_days": 5, "has_ssl": True, "reputation_score": 0.05}], "label": "QUARANTINE", "difficulty": 3 }, { "sender": "it-support@manipal-helpdesk.com", # Look-alike domain "subject": "Mandatory Password Reset", "body": "As per the new MIT security policy, all students must reset their password today.", "raw_headers": "Received: from mail-delivery.online... SPF: pass; DKIM: pass;", # Attacker set up SPF/DKIM correctly! "auth_results": {"SPF": "pass", "DKIM": "pass", "DMARC": "none"}, "urls": [{"url": "http://manipal-helpdesk.com/reset", "is_shortened": False, "domain_age_days": 1, "has_ssl": False, "reputation_score": 0.1}], "label": "QUARANTINE", "difficulty": 3 }, { "sender": "prof.sharma@mit-manipal.org", # Wrong TLD "subject": "Final Exam Paper Leak?", "body": "I suspect the paper has leaked. Look at this screenshot and confirm if these are your questions.", "raw_headers": "Received: from sendgrid.net... SPF: pass;", "auth_results": {"SPF": "pass", "DKIM": "none", "DMARC": "none"}, "urls": [{"url": "https://dropbox-files.com/s/xyz", "is_shortened": False, "domain_age_days": 4, "has_ssl": True, "reputation_score": 0.2}], "label": "QUARANTINE", "difficulty": 3 }, { "sender": "accounts@google-security.info", "subject": "Critical Security Alert", "body": "Someone just used your password to try to sign in to your account. Go to your Google account now.", "raw_headers": "Received: from host-12.xyz... SPF: fail;", "auth_results": {"SPF": "fail", "DKIM": "none", "DMARC": "fail"}, "urls": [{"url": "https://google-secure-login.info", "is_shortened": False, "domain_age_days": 2, "has_ssl": True, "reputation_score": 0.01}], "label": "QUARANTINE", "difficulty": 3 }, { "sender": "library@manipal.edu", "subject": "Overdue Book Notice", "body": "Your copy of 'Computer Networks' is overdue. Click to pay the fine of ₹50.", "raw_headers": "Received: from mail.manipal.edu... SPF: pass; DKIM: pass;", "auth_results": {"SPF": "pass", "DKIM": "pass", "DMARC": "pass"}, "urls": [{"url": "https://portal.manipal.edu/pay", "is_shortened": False, "domain_age_days": 4000, "has_ssl": True, "reputation_score": 1.0}], "label": "INBOX", "difficulty": 3 # Difficult because it looks like a phishing lure but is legit. } ] return base_data async def reset(self): self.current_step = 0 return self._get_result() def _get_result(self, reward=0.0, done=False): if self.current_step >= len(self.dataset): obs = MyEnvV4Observation( sender="N/A", subject="N/A", body="N/A", raw_headers="", hop_count=0, auth_results={}, urls=[], echoed_message="End of Session" ) return SimpleNamespace(observation=obs, reward=reward, done=True) data = self.dataset[self.current_step] obs = MyEnvV4Observation( sender=data["sender"], subject=data["subject"], body=data["body"], raw_headers=data["raw_headers"], hop_count=data["raw_headers"].count("Received:"), auth_results=data["auth_results"], urls=[URLInfo(**u) for u in data["urls"]], echoed_message=f"Task {self.current_step + 1}/{len(self.dataset)}" ) return SimpleNamespace(observation=obs, reward=reward, done=done) async def step(self, action: MyEnvV4Action): if self.current_step >= len(self.dataset): return self._get_result(done=True) target = self.dataset[self.current_step] correct = target["label"] prediction = action.message.strip().upper() # SOPHISTICATED REWARD LOGIC reward = 0.0 if prediction == correct: # Perfect Match: reward scales with difficulty reward = 1.0 + (target["difficulty"] * 0.1) elif correct in ["SPAM", "QUARANTINE"] and prediction in ["SPAM", "QUARANTINE"]: # Partial Credit: Recognized threat but misclassified type reward = 0.4 elif correct == "QUARANTINE" and prediction == "INBOX": # Dangerous Failure: Penalty for letting a threat into the Inbox reward = -1.5 elif correct == "INBOX" and prediction == "QUARANTINE": # False Positive: Penalty for blocking legitimate mail reward = -0.5 # Add Reasoning Bonus (Explainability) if hasattr(action, 'reasoning') and action.reasoning and len(action.reasoning) > 30: # Small bonus if agent provides a justification reward += 0.05 # Normalize reward to [0, 1] range as per OpenEnv specs (clipping/rescaling) # However, many environments allow negative for penalties; we clip to [0,1] for final score final_reward = max(0.0, min(1.0, reward)) self.current_step += 1 done = self.current_step >= len(self.dataset) return self._get_result(reward=final_reward, done=done) async def state(self): return {"current_step": self.current_step, "total_tasks": len(self.dataset)} # Global instance for the server my_env = MyEnvV4Env() app = FastAPI() @app.get("/reset") async def reset(): res = await my_env.reset() return {"observation": res.observation, "reward": res.reward, "done": res.done} @app.post("/step") async def step(action: MyEnvV4Action): res = await my_env.step(action) return {"observation": res.observation, "reward": res.reward, "done": res.done} @app.get("/state") async def state(): return await my_env.state()