Spaces:
Sleeping
Sleeping
File size: 8,799 Bytes
45dfcaa | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 | """
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
|