mohdbelal010/SecureAI-Gaurd

1import logging2import os3from typing import Optional4 5logger = logging.getLogger(__name__)6 7# Risk keywords weighted by severity8HIGH_RISK_KEYWORDS = [9    "click here", "verify your account", "suspend", "suspended", "immediate action",10    "you've won", "claim now", "download", "install", "executable", "wire transfer",11    "urgent", "emergency", "credentials", "password expired", "password has expired",12    "customs fee", "stranded", "send money", "limited time", "pre-approved",13    "reset here", "action required", "confirm your identity",14]15 16MEDIUM_RISK_KEYWORDS = [17    "click", "verify", "account", "update", "confirm", "link", "http",18    "invoice", "document", "security", "alert", "warning", "expire",19    "offer", "free", "win", "prize", "loan", "approved",20]21 22# Safe-content patterns: messages matching these are very likely legitimate23SAFE_PATTERNS = [24    "meeting confirmed", "has shipped", "expected delivery",25    "get back to you", "dinner tonight", "appointment",26    "monthly statement", "conference room", "thanks for",27    "don't forget", "reminder:", "order #", "shipped",28    "end of day", "see you", "good morning", "good evening",29    "happy birthday", "thank you",30]31 32 33class HFRiskScorer:34    """35    HuggingFace-backed risk scorer with keyword-primary approach.36 37    The HF model (distilbert-sst2) is a *sentiment* classifier, so it can38    mis-score benign messages (e.g. "dentist appointment" → NEGATIVE sentiment39    → high "risk").  To prevent false positives the scorer uses keyword40    heuristics as the **primary** signal (80 %) and blends the HF model41    score in only as a secondary correction (20 %).42    """43 44    def __init__(self):45        self.classifier = None46        self._try_load_model()47 48    def _try_load_model(self):49        hf_token = os.environ.get("HF_TOKEN", "")50        try:51            from transformers import pipeline52            import torch53 54            model_name = os.environ.get(55                "HF_RISK_MODEL", "distilbert-base-uncased-finetuned-sst-2-english"56            )57            device = 0 if (hasattr(torch, "cuda") and torch.cuda.is_available()) else -158            self.classifier = pipeline(59                "text-classification", model=model_name, device=device,60                token=hf_token if hf_token else None,61            )62            logger.info("HF model loaded: %s", model_name)63        except Exception as exc:64            logger.warning("HF model unavailable (%s). Using keyword fallback.", exc)65            self.classifier = None66 67    def score_text(self, text: str) -> float:68        """Return risk score in [0.0, 1.0].69 70        Uses keyword heuristics as the primary signal (80 %) and the HF71        sentiment model only as a weak secondary signal (20 %).72        """73        kw_score = self._keyword_score(text)74 75        hf_score = kw_score  # default: same as keyword if model unavailable76        if self.classifier is not None:77            try:78                result = self.classifier(text[:512])[0]79                label = result["label"].upper()80                score = float(result["score"])81                # NEGATIVE / LABEL_0 → risky; POSITIVE / LABEL_1 → safe82                if label in ("NEGATIVE", "LABEL_0"):83                    hf_score = min(score, 1.0)84                else:85                    hf_score = max(1.0 - score, 0.0)86            except Exception as exc:87                logger.warning("Classifier error: %s", exc)88 89        # Blend: keyword-primary (80%) + HF-secondary (20%)90        blended = 0.80 * kw_score + 0.20 * hf_score91        return round(min(max(blended, 0.0), 1.0), 4)92 93    def _keyword_score(self, text: str) -> float:94        t = text.lower()95 96        # Check for safe-content patterns first — if the message looks97        # clearly legitimate, give it a very low risk score.98        safe_hits = sum(1 for p in SAFE_PATTERNS if p in t)99        if safe_hits >= 1:100            high = sum(1 for kw in HIGH_RISK_KEYWORDS if kw in t)101            # Only override if there are no high-risk keywords present102            if high == 0:103                return round(max(0.05 - safe_hits * 0.01, 0.0), 4)104 105        high = sum(1 for kw in HIGH_RISK_KEYWORDS if kw in t)106        medium = sum(1 for kw in MEDIUM_RISK_KEYWORDS if kw in t)107        score = min(high * 0.25 + medium * 0.08, 1.0)108        return round(score, 4)109