Version_5/src/sub_agents.py

import os
from groq import Groq
from .schemas import AgentReport

# Initialize the Groq Cloud Client
client = Groq(api_key=os.environ.get("GROQ_API_KEY"))

# ── SYSTEM OVERRIDE 1: AMTSO & Simulated Tests ──
SIMULATION_OVERRIDE = """
CRITICAL RULE - SIMULATED THREATS: 
If the text, HTML, or content explicitly claims to be a "Test Page", "Dummy Page", "AMTSO Phishing Check", or states it is designed to test anti-malware/antivirus software, YOU MUST treat it as a real threat. Return a claim of "PHISHING" with a high confidence score, and explicitly state in the evidence array that it is a simulated testing environment. Do not mark it as LEGITIMATE just because it is a benign test.
"""

# ── SYSTEM OVERRIDE 2: Legitimate Brands & Spoofing ──
TRUSTED_BRAND_OVERRIDE = """
CRITICAL RULE - LEGITIMATE BRANDS & LINK SPOOFING:
1. If the ACTUAL DESTINATION URL belongs to a verified, universally trusted root domain (e.g., 'github.com', 'google.com') with no typosquatting, you MUST return "LEGITIMATE".
2. LINK SPOOFING FATAL FLAW: If the visible text or email claims to be a trusted brand, but the underlying hyperlink (href) or actual destination URL points to a completely different, unrelated domain (e.g., text says 'github.com' but URL is 'login-update-security.com'), YOU MUST FLAG IT AS "PHISHING" with 1.0 confidence.
"""

def run_agent(role_prompt: str, target_data: str) -> AgentReport:
    # We call Groq's hosted 8B infrastructure instead of local resources
    response = client.chat.completions.create(
        model="llama-3.1-8b-instant",
        messages=[
            {"role": "system", "content": role_prompt},
            {"role": "user", "content": target_data}
        ],
        # Groq enforces JSON mode via this parameter block
        response_format={"type": "json_object"},
        temperature=0.0
    )
    return AgentReport.model_validate_json(response.choices[0].message.content)

def agent_url_analyst(url: str) -> AgentReport:
    prompt = (
        "You are a URL Shadows Analyst. Respond strictly in valid JSON matching this schema:\n"
        "{\"claim\": \"PHISHING\"|\"LEGITIMATE\"|\"AMBIGUOUS\", \"confidence\": float, \"evidence\": [string]}\n"
        "Analyze the structure for structural entropy, unusual domains, and brand keywords.\n"
        + TRUSTED_BRAND_OVERRIDE
    )
    return run_agent(prompt, f"URL to analyze: {url}")

def agent_html_structure(dom_json: str) -> AgentReport:
    prompt = (
        "You are an HTML Code Analyst. Respond strictly in valid JSON matching this schema:\n"
        "{\"claim\": \"PHISHING\"|\"LEGITIMATE\"|\"AMBIGUOUS\", \"confidence\": float, \"evidence\": [string]}\n"
        "Analyze these DOM elements. Flag external tracking inputs or forms posting to alternative servers.\n"
        + SIMULATION_OVERRIDE
    )
    return run_agent(prompt, f"DOM Data: {dom_json}")

def agent_content_semantics(email_body: str) -> AgentReport:
    prompt = (
        "You are a Phishing Copywriter Critic. Respond strictly in valid JSON matching this schema:\n"
        "{\"claim\": \"PHISHING\"|\"LEGITIMATE\"|\"AMBIGUOUS\", \"confidence\": float, \"evidence\": [string]}\n"
        "Extract semantic compliance anomalies, high emotional coercion markers, and urgency loops.\n"
        + SIMULATION_OVERRIDE
    )
    return run_agent(prompt, f"Email Body: {email_body}")

def agent_brand_impersonation(email_body: str, sender: str) -> AgentReport:
    prompt = (
        "You are an Identity Protection Agent. Respond strictly in valid JSON matching this schema:\n"
        "{\"claim\": \"PHISHING\"|\"LEGITIMATE\"|\"AMBIGUOUS\", \"confidence\": float, \"evidence\": [string]}\n"
        "Flag mismatched target namespaces where sender domains do not match corporate identifiers.\n"
        + SIMULATION_OVERRIDE
        + "\n"
        + TRUSTED_BRAND_OVERRIDE
    )
    return run_agent(prompt, f"Sender: {sender}\nBody: {email_body}")