categorization.py

"""
Risk categorization module for phishing detection results.

This module provides functions to categorize phishing probability scores
into risk categories and binary classifications.
"""

from enum import Enum
from typing import Tuple


class RiskCategory(str, Enum):
    """Risk category based on phishing probability score."""
    SAFE = "Safe"
    LOW = "Low"
    MODERATE = "Moderate"
    HIGH = "High"
    CRITICAL = "Critical"


class BinaryClassification(str, Enum):
    """Binary classification of phishing detection result."""
    LEGITIMATE = "Legitimate"
    PHISHING = "Phishing"


# Risk category thresholds (score is 0-100 scale)
RISK_THRESHOLDS = {
    RiskCategory.SAFE: (0, 30),        # 0-30%: score <= 30
    RiskCategory.LOW: (30, 50),        # 31-50%: 30 < score <= 50
    RiskCategory.MODERATE: (50, 70),   # 51-70%: 50 < score <= 70
    RiskCategory.HIGH: (70, 85),       # 71-85%: 70 < score <= 85
    RiskCategory.CRITICAL: (85, 101),  # 86-100%: score > 85
}

# Binary classification threshold
PHISHING_THRESHOLD = 70  # score > 70 is classified as Phishing


def get_risk_category(phish_probability_score: float) -> RiskCategory:
    """
    Determine the risk category based on phishing probability score.

    Thresholds:
    - Safe: 0-30%
    - Low: 31-50%
    - Moderate: 51-70%
    - High: 71-85%
    - Critical: 86-100%

    Args:
        phish_probability_score: Phishing probability score (0-100 scale)

    Returns:
        RiskCategory: The corresponding risk category
    """
    if phish_probability_score <= 30:
        return RiskCategory.SAFE
    elif phish_probability_score <= 50:
        return RiskCategory.LOW
    elif phish_probability_score <= 70:
        return RiskCategory.MODERATE
    elif phish_probability_score <= 85:
        return RiskCategory.HIGH
    else:
        return RiskCategory.CRITICAL


def get_binary_classification(phish_probability_score: float) -> BinaryClassification:
    """
    Determine the binary classification based on phishing probability score.

    Args:
        phish_probability_score: Phishing probability score (0-100 scale)

    Returns:
        BinaryClassification: Legitimate if score <= 70, Phishing otherwise
    """
    if phish_probability_score <= PHISHING_THRESHOLD:
        return BinaryClassification.LEGITIMATE
    else:
        return BinaryClassification.PHISHING


def categorize_phishing_result(phish_probability: float) -> Tuple[RiskCategory, BinaryClassification, float]:
    """
    Categorize a phishing detection result.
    
    This function takes a phishing probability (0-1 scale) and returns:
    - Risk category (Safe, Low, Moderate, Dangerous, Critical)
    - Binary classification (Legitimate or Phishing)
    - The probability score on a 0-100 scale
    
    Args:
        phish_probability: Phishing probability from the model (0-1 scale)
    
    Returns:
        Tuple containing:
        - RiskCategory: The risk category
        - BinaryClassification: The binary classification
        - float: The probability score on 0-100 scale
    """
    # Convert from 0-1 scale to 0-100 scale
    score_100 = phish_probability * 100
    
    risk_category = get_risk_category(score_100)
    binary_classification = get_binary_classification(score_100)
    
    return risk_category, binary_classification, score_100