Spaces:
Sleeping
Sleeping
| """ | |
| Phishing detection inference pipeline. | |
| This module provides a complete inference pipeline for phishing detection including: | |
| - Model loading from Hopsworks Model Registry | |
| - Feature preprocessing | |
| - Prediction with confidence scores | |
| - Optional end-to-end URL scanning and prediction | |
| """ | |
| import os | |
| import logging | |
| from typing import Tuple, Optional, Dict, Any | |
| import pandas as pd | |
| import joblib | |
| from phising_detection.utils.hopsworks_utils import connect_to_hopsworks | |
| from phising_detection.models.model_utils.data_prep import FEATURE_COLUMNS, CONTINUOUS_FEATURES | |
| from phising_detection.features.urlscan_features import extract_features | |
| from phising_detection.utils.urlscan import URLScanClient | |
| logger = logging.getLogger(__name__) | |
| class PhishingDetectionPipeline: | |
| """ | |
| End-to-end inference pipeline for phishing detection. | |
| This class handles: | |
| - Loading trained models from Hopsworks | |
| - Preprocessing extracted features | |
| - Running inference | |
| - Optionally scanning URLs with URLScan.io | |
| """ | |
| def __init__( | |
| self, | |
| model_name: str = "phishing_detector", | |
| model_version: Optional[int] = None, | |
| urlscan_api_key: Optional[str] = None | |
| ): | |
| """ | |
| Initialize the inference pipeline. | |
| Args: | |
| model_name: Name of the model in Hopsworks Model Registry | |
| model_version: Specific version to load (None = latest) | |
| urlscan_api_key: URLScan.io API key (optional, only needed for URL scanning) | |
| """ | |
| self.model_name = model_name | |
| self.model_version = model_version | |
| self.model = None | |
| self.scaler = None | |
| self.feature_names = None | |
| self.urlscan_client = None | |
| if urlscan_api_key: | |
| self.urlscan_client = URLScanClient(api_key=urlscan_api_key) | |
| def load_model_from_hopsworks(self) -> None: | |
| """ | |
| Load model, scaler, and feature names from Hopsworks Model Registry. | |
| Raises: | |
| Exception: If model loading fails | |
| """ | |
| logger.info(f"Loading model '{self.model_name}' from Hopsworks Model Registry...") | |
| # Connect to Hopsworks | |
| project = connect_to_hopsworks() | |
| mr = project.get_model_registry() | |
| # Get model from registry | |
| if self.model_version: | |
| model_registry = mr.get_model(self.model_name, version=self.model_version) | |
| else: | |
| model_registry = mr.get_model(self.model_name) | |
| logger.info(f"Found model: {self.model_name} version {model_registry.version}") | |
| # Download model artifacts to a temporary directory | |
| model_dir = model_registry.download() | |
| logger.info(f"Model artifacts downloaded to: {model_dir}") | |
| # Load model | |
| model_path = os.path.join(model_dir, "model.pkl") | |
| self.model = joblib.load(model_path) | |
| logger.info(f"Model loaded from {model_path}") | |
| # Load scaler | |
| scaler_path = os.path.join(model_dir, "scaler.pkl") | |
| self.scaler = joblib.load(scaler_path) | |
| logger.info(f"Scaler loaded from {scaler_path}") | |
| # Load feature names | |
| feature_names_path = os.path.join(model_dir, "feature_names.txt") | |
| with open(feature_names_path, 'r') as f: | |
| self.feature_names = [line.strip() for line in f.readlines()] | |
| logger.info(f"Feature names loaded: {self.feature_names}") | |
| logger.info("Model pipeline initialized successfully!") | |
| def preprocess_features(self, features_dict: Dict[str, Any]) -> pd.DataFrame: | |
| """ | |
| Preprocess extracted features to match model input format. | |
| Args: | |
| features_dict: Dictionary of extracted features | |
| Returns: | |
| Preprocessed DataFrame ready for model inference | |
| Raises: | |
| ValueError: If model is not loaded | |
| """ | |
| if self.model is None or self.scaler is None or self.feature_names is None: | |
| raise ValueError("Model not loaded. Call load_model_from_hopsworks() first.") | |
| # Create DataFrame with a single row | |
| df = pd.DataFrame([features_dict]) | |
| # Ensure all required features are present | |
| for feature in self.feature_names: | |
| if feature not in df.columns: | |
| logger.warning(f"Feature '{feature}' missing, filling with default value") | |
| df[feature] = 0 | |
| # Select and order features to match training | |
| df = df[self.feature_names] | |
| # Handle missing values (fill with median approximation) | |
| df = df.fillna(df.median()) | |
| # Apply scaling to continuous features | |
| df_scaled = df.copy() | |
| df_scaled[CONTINUOUS_FEATURES] = self.scaler.transform(df[CONTINUOUS_FEATURES]) | |
| return df_scaled | |
| def predict(self, features_dict: Dict[str, Any]) -> Dict[str, Any]: | |
| """ | |
| Run inference on extracted features. | |
| Args: | |
| features_dict: Dictionary of extracted features | |
| Returns: | |
| Dictionary containing prediction results: | |
| - prediction: "PHISHING" or "LEGITIMATE" | |
| - confidence: Confidence score (0-1) | |
| - phishing_probability: Probability of phishing (0-1) | |
| - legitimate_probability: Probability of legitimate (0-1) | |
| - is_phishing: Boolean flag | |
| Raises: | |
| ValueError: If model is not loaded | |
| """ | |
| if self.model is None: | |
| raise ValueError("Model not loaded. Call load_model_from_hopsworks() first.") | |
| # Preprocess features | |
| X = self.preprocess_features(features_dict) | |
| # Run inference | |
| logger.info("Running model inference...") | |
| prediction_proba = self.model.predict_proba(X)[0] | |
| prediction = self.model.predict(X)[0] | |
| # prediction: 0 = legitimate, 1 = phishing | |
| is_phishing = bool(prediction) | |
| confidence = prediction_proba[1] if is_phishing else prediction_proba[0] | |
| result = { | |
| "prediction": "PHISHING" if is_phishing else "LEGITIMATE", | |
| "confidence": float(confidence), | |
| "phishing_probability": float(prediction_proba[1]), | |
| "legitimate_probability": float(prediction_proba[0]), | |
| "is_phishing": is_phishing | |
| } | |
| logger.info(f"Prediction: {result['prediction']} (confidence: {result['confidence']:.4f})") | |
| return result | |
| def predict_url(self, url: str) -> Dict[str, Any]: | |
| """ | |
| End-to-end prediction: scan URL and predict if it's phishing. | |
| Args: | |
| url: URL to analyze | |
| Returns: | |
| Dictionary containing: | |
| - prediction: "PHISHING" or "LEGITIMATE" | |
| - confidence: Confidence score (0-1) | |
| - phishing_probability: Probability of phishing (0-1) | |
| - legitimate_probability: Probability of legitimate (0-1) | |
| - is_phishing: Boolean flag | |
| - features: Dictionary of extracted features | |
| - scan_uuid: URLScan UUID (if available) | |
| - error: Error message (if any) | |
| Raises: | |
| ValueError: If URLScan client is not initialized | |
| """ | |
| if self.urlscan_client is None: | |
| raise ValueError("URLScan client not initialized. Provide urlscan_api_key in constructor.") | |
| try: | |
| # Step 1: Submit URL to URLScan.io and wait for results | |
| logger.info(f"Scanning URL: {url}") | |
| scan_result = self.urlscan_client.submit_and_wait(url) | |
| if not scan_result: | |
| return { | |
| "error": "Failed to get scan results from URLScan.io", | |
| "prediction": "ERROR", | |
| "confidence": 0.0 | |
| } | |
| # Step 2: Extract features from scan results | |
| logger.info("Extracting features from scan results...") | |
| features = extract_features(scan_result) | |
| logger.info(f"Extracted features: {features}") | |
| # Step 3: Run inference | |
| prediction_result = self.predict(features) | |
| # Add additional information | |
| prediction_result["features"] = features | |
| prediction_result["scan_uuid"] = scan_result.get("task", {}).get("uuid", "N/A") | |
| return prediction_result | |
| except Exception as e: | |
| logger.error(f"Error during prediction: {e}", exc_info=True) | |
| return { | |
| "error": str(e), | |
| "prediction": "ERROR", | |
| "confidence": 0.0 | |
| } | |
| def is_loaded(self) -> bool: | |
| """ | |
| Check if the model is loaded and ready for inference. | |
| Returns: | |
| True if model is loaded, False otherwise | |
| """ | |
| return self.model is not None and self.scaler is not None and self.feature_names is not None | |