heuristics.py

import re
from typing import Dict, List, Optional, Tuple
import pathlib
from pathlib import Path

FILE_PATH = Path("/shield-intel/disposable-email-domains.txt")

def load_disposable_domains(file_path: pathlib.Path) -> set[str]:
    """
    Read the given file, ignore blank lines and whitespace,
    and return a set of domain strings.
    """
    domains = set()
    with file_path.open("r", encoding="utf-8") as f:
        for line in f:
            domain = line.strip()
            if domain:
                domains.add(domain.lower())
    return domains

DISPOSABLE_DOMAINS = load_disposable_domains(FILE_PATH)
SUSPICIOUS_TLDS = {".tk", ".ml", ".ga", ".cf", ".gq", ".loan", ".work", ".date"}
KNOWN_ABUSE_IPS: set = set()

COMMON_SPAM_PATTERNS = [
    r"(?i)buy\s+now.*click\s+here",
    r"(?i)limited\s+offer.*expires",
    r"(?i)congratulations.*won",
    r"(?i)click\s+the\s+link.*verify",
    r"(?i)account.*suspended.*click",
    r"(?i)confirm.*password.*urgent",
    r"(?i)free\s+access.*credit",
]

COMMON_PROMPT_ATTACK_PATTERNS = [
    r"(?i)ignore\s+(all\s+)?(previous|above|prior)\s+instructions",
    r"(?i)disregard\s+(all\s+)?(previous|above)\s+(instructions|directives|rules)",
    r"(?i)you\s+are\s+(now|an?\s+free)\s+",
    r"(?i)system\s+prompt",
    r"(?i)forget\s+(everything|all\s+previous)",
    r"(?i)new\s+instruction",
    r"(?i)act\s+as\s+(if\s+you\s+are|an?\s+AI\s+with)",
    r"(?i)jailbreak",
    r"(?i)do\s+not\s+follow\s+the\s+(rules|guidelines|restrictions)",
    r"(?i)you\s+have\s+no\s+(restrictions|limitations|boundaries)",
]


def check_disposable_email(email: str) -> Optional[str]:
    if not email or "@" not in email:
        return None
    domain = email.split("@", 1)[1].lower().strip()
    if domain in DISPOSABLE_DOMAINS:
        return "Disposable email detected"
    tld = "." + domain.rsplit(".", 1)[-1]
    if tld in SUSPICIOUS_TLDS:
        return f"Suspicious email TLD ({tld})"
    return None


def check_known_abuse_ip(ip: str) -> Optional[str]:
    if not ip:
        return None
    if ip in KNOWN_ABUSE_IPS:
        return "IP linked to previous abuse"
    return None


def check_spam_content(content: str) -> List[str]:
    if not content:
        return []
    matches = []
    for pattern in COMMON_SPAM_PATTERNS:
        if re.search(pattern, content):
            matches.append("Content resembles spam pattern")
            break
    return matches


def check_prompt_attack(content: str) -> List[str]:
    if not content:
        return []
    matches = []
    for pattern in COMMON_PROMPT_ATTACK_PATTERNS:
        if re.search(pattern, content):
            matches.append("Prompt injection attempt detected")
            break
    return matches


def check_high_velocity_signup(signup_time: Optional[str]) -> List[str]:
    if not signup_time:
        return []
    return []


def run_heuristics(
    email: Optional[str] = None,
    ip: Optional[str] = None,
    content: Optional[str] = None,
) -> Tuple[int, List[str]]:
    reasons: List[str] = []
    score = 0

    email_result = check_disposable_email(email)
    if email_result:
        reasons.append(email_result)
        score += 25

    ip_result = check_known_abuse_ip(ip)
    if ip_result:
        reasons.append(ip_result)
        score += 30

    spam_reasons = check_spam_content(content)
    reasons.extend(spam_reasons)
    if spam_reasons:
        score += 20

    attack_reasons = check_prompt_attack(content)
    reasons.extend(attack_reasons)
    if attack_reasons:
        score += 35

    return min(score, 100), reasons