Spaces:
Sleeping
Sleeping
File size: 3,931 Bytes
f9af852 79f8ec5 f9af852 79f8ec5 f9af852 79f8ec5 f9af852 79f8ec5 f9af852 92f63a9 | 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 69 70 71 | 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}") |