screenshot_hasher.py

# ============================================================
# PhishGuard AI - cnn/screenshot_hasher.py
# Perceptual hash-based brand impersonation detector.
#
# Compares webpage screenshots against reference hashes of
# known brand login pages using imagehash.phash.
#
# brand_boost = 0.25 if hamming_distance < 10 else 0.0
# ============================================================

from __future__ import annotations

import io
import json
import logging
from pathlib import Path
from typing import Tuple, Optional, Dict, List

from PIL import Image

logger = logging.getLogger("phishguard.cnn.hasher")

# ── Try to use imagehash, fall back to custom implementation ─────────
_imagehash_available = False
try:
    import imagehash
    _imagehash_available = True
except ImportError:
    logger.info("imagehash not installed — using built-in phash")

HASH_DB_PATH = Path(__file__).parent / "brand_hashes.json"


class BrandHashDetector:
    """
    Perceptual hash-based brand impersonation detector.
    Compares screenshots against reference hashes of 10 major brands.
    """

    BRANDS: List[str] = [
        "paypal", "google", "apple", "microsoft", "amazon",
        "chase", "netflix", "facebook", "instagram", "wellsfargo",
    ]

    BRAND_DOMAINS: Dict[str, str] = {
        "paypal": "paypal.com",
        "google": "google.com",
        "apple": "apple.com",
        "microsoft": "microsoft.com",
        "amazon": "amazon.com",
        "chase": "chase.com",
        "netflix": "netflix.com",
        "facebook": "facebook.com",
        "instagram": "instagram.com",
        "wellsfargo": "wellsfargo.com",
    }

    def __init__(self, hash_db_path: Optional[Path] = None) -> None:
        self._hash_db_path = hash_db_path or HASH_DB_PATH
        self._reference_hashes: Dict[str, dict] = {}
        self._load_reference_hashes()

    def _load_reference_hashes(self) -> None:
        """Load reference hashes from JSON database."""
        if self._hash_db_path.exists():
            try:
                with open(self._hash_db_path) as f:
                    self._reference_hashes = json.load(f)
                logger.info(f"Loaded {len(self._reference_hashes)} brand hashes")
            except Exception as e:
                logger.warning(f"Failed to load brand hashes: {e}")
                self._reference_hashes = {}
        else:
            logger.info("No brand hash DB found — brand detection disabled")
            self._reference_hashes = {}

    def compute_hash(self, img_bytes: bytes, hash_size: int = 16) -> Optional[int]:
        """
        Compute perceptual hash of an image.
        Uses imagehash.phash if available, otherwise custom DCT-less implementation.
        """
        try:
            img = Image.open(io.BytesIO(img_bytes))

            if _imagehash_available:
                h = imagehash.phash(img, hash_size=hash_size)
                return int(str(h), 16)
            else:
                return self._custom_phash(img, hash_size)

        except Exception as e:
            logger.warning(f"Hash computation failed: {e}")
            return None

    def _custom_phash(self, img: Image.Image, hash_size: int = 16) -> int:
        """Fallback perceptual hash (mean-based, no DCT)."""
        img = img.convert("L").resize((hash_size, hash_size), Image.LANCZOS)
        pixels = list(img.getdata())
        avg = sum(pixels) / len(pixels)
        bits = "".join("1" if p > avg else "0" for p in pixels)
        return int(bits, 2)

    def hamming_distance(self, h1: int, h2: int) -> int:
        """Count bit differences between two hashes. 0 = identical."""
        return bin(h1 ^ h2).count("1")

    def detect(
        self,
        screenshot_bytes: bytes,
        url: str = "",
        threshold: int = 10,
    ) -> Tuple[bool, str, float]:
        """
        Detect brand impersonation from screenshot.

        Returns:
            (is_impersonation, brand_name, confidence)
            is_impersonation: True if page looks like a brand but URL doesn't match
            brand_name: detected brand name or ""
            confidence: 0.0-1.0 similarity score
        """
        page_hash = self.compute_hash(screenshot_bytes)
        if page_hash is None:
            return False, "", 0.0

        url_lower = url.lower()
        best_match: Optional[str] = None
        best_distance = 999
        best_confidence = 0.0

        for brand, entry in self._reference_hashes.items():
            try:
                stored_hash = int(entry["hash"])
                distance = self.hamming_distance(page_hash, stored_hash)
                confidence = max(0.0, 1.0 - distance / 256.0)

                if distance < best_distance:
                    best_distance = distance
                    best_match = brand
                    best_confidence = confidence

            except (ValueError, KeyError):
                continue

        if best_match and best_distance <= threshold:
            legit_domain = self.BRAND_DOMAINS.get(best_match, f"{best_match}.com")

            # Check if URL belongs to legitimate domain
            if legit_domain not in url_lower:
                return True, best_match, best_confidence
            else:
                return False, best_match, best_confidence

        return False, "", 0.0

    def register_brand(
        self,
        brand_name: str,
        domain: str,
        screenshot_bytes: bytes,
    ) -> bool:
        """Register a brand's reference screenshot hash."""
        h = self.compute_hash(screenshot_bytes)
        if h is None:
            return False

        self._reference_hashes[brand_name] = {
            "domain": domain,
            "hash": str(h),
        }

        # Save to disk
        try:
            with open(self._hash_db_path, "w") as f:
                json.dump(self._reference_hashes, f, indent=2)
            logger.info(f"Registered brand: {brand_name} ({domain})")
            return True
        except Exception as e:
            logger.error(f"Failed to save brand hash: {e}")
            return False


# ── Legacy compatibility ─────────────────────────────────────────────
_detector = BrandHashDetector()


def check_brand_impersonation(
    screenshot_bytes: bytes,
    url: str,
    similarity_threshold: int = 10,
) -> dict:
    """Legacy wrapper for backward compatibility."""
    is_impersonation, brand, confidence = _detector.detect(
        screenshot_bytes, url, similarity_threshold,
    )

    if is_impersonation:
        return {
            "impersonation_detected": True,
            "impersonated_brand": brand,
            "legitimate_domain": _detector.BRAND_DOMAINS.get(brand, ""),
            "visual_similarity": round(confidence, 3),
        }
    elif brand:
        return {
            "impersonation_detected": False,
            "matched_brand": brand,
            "note": "legitimate site",
        }
    else:
        return {
            "impersonation_detected": False,
            "reason": "no_brand_match",
        }