Update env.py

env.py

@@ -1,52 +1,191 @@
 import asyncio
-from typing import Optional, List
+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
+from models import MyEnvV4Observation, MyEnvV4Action, URLInfo
 
 class MyEnvV4Env(Environment):
     def __init__(self):
         super().__init__()
-        # Realistic Dataset with Digital Seduction/Phishing markers
-        self.dataset = [
+        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": "dean.office@manipal.edu",
-                "subject": "B.Tech Lab Exam Schedule",
-                "body": "Please find the attached PDF for the upcoming CSE lab exams.",
-                "headers": ["SPF: Pass", "DKIM: Pass"],
-                "label": "INBOX"
+                "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": "verify-account@security-amazon.net",
-                "subject": "Urgent: Your account is locked!",
-                "body": "Digital Seduction Alert: High urgency used. Click http://bit.ly/fake-link to unlock.",
-                "headers": ["SPF: Fail", "DMARC: Fail"],
-                "label": "QUARANTINE"
+                "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": "prize@lottery-winner.co",
-                "subject": "Congratulations! You won $10,000",
-                "body": "Reply with your bank details to claim your cash prize immediately.",
-                "headers": ["SPF: Neutral"],
-                "label": "SPAM"
+                "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.
             }
         ]
-        self.current_step = 0
+        return base_data
 
     async def reset(self):
         self.current_step = 0
         return self._get_result()
-        
-    async def state(self):
-        """Returns the current state of the environment."""
-        return {"current_step": self.current_step, "total_steps": len(self.dataset)}
 
     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", 
-                headers=[], echoed_message="End of Data"
+                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)
         
@@ -55,8 +194,11 @@ class MyEnvV4Env(Environment):
             sender=data["sender"],
             subject=data["subject"],
             body=data["body"],
-            headers=data["headers"],
-            echoed_message=f"Step {self.current_step + 1}"
+            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)
 
@@ -64,40 +206,56 @@ class MyEnvV4Env(Environment):
         if self.current_step >= len(self.dataset):
             return self._get_result(done=True)
 
-        correct_label = self.dataset[self.current_step]["label"]
-        # Exact match reward logic for 0.0 - 1.0 range
-        reward = 1.0 if action.message.strip().upper() == correct_label else 0.0
+        target = self.dataset[self.current_step]
+        correct = target["label"]
+        prediction = action.message.strip().upper()
         
-        self.current_step += 1
-        done = self.current_step >= len(self.dataset)
+        # SOPHISTICATED REWARD LOGIC
+        reward = 0.0
         
-        return self._get_result(reward=reward, done=done)
-
-    async def close(self):
-        pass
+        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
 
-    @classmethod
-    async def from_docker_image(cls, image_name: str):
-        """Simulated helper for local/containerized runs."""
-        return cls()
+        # 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)}
 
-# Create the environment instance
+# Global instance for the server
 my_env = MyEnvV4Env()
-
-# defining the FastAPI app and routes
 app = FastAPI()
 
 @app.get("/reset")
 async def reset():
-    result = await my_env.reset()
-    return {"observation": result.observation, "reward": result.reward, "done": result.done}
+    res = await my_env.reset()
+    return {"observation": res.observation, "reward": res.reward, "done": res.done}
 
 @app.post("/step")
 async def step(action: MyEnvV4Action):
-    result = await my_env.step(action)
-    return {"observation": result.observation, "reward": result.reward, "done": result.done}
+    res = await my_env.step(action)
+    return {"observation": res.observation, "reward": res.reward, "done": res.done}
 
 @app.get("/state")
 async def state():