env.py

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()