Spaces:
Paused
Paused
| import pandas as pd | |
| import xgboost as xgb | |
| import logging | |
| import json | |
| from pathlib import Path | |
| from datetime import datetime | |
| from typing import Dict, Any | |
| logger = logging.getLogger(__name__) | |
| class XGBoostClassifier: | |
| """ | |
| Module 3: Attack Classification. | |
| Supervised XGBoost model to classify authentication events into specific attack categories. | |
| """ | |
| def __init__(self, model_dir: str, config: Dict[str, Any]): | |
| self.model_dir = Path(model_dir) | |
| self.model = None | |
| self.config = config | |
| self.version = "1.0.0" | |
| self.labels = self.config.get("classification_labels", [ | |
| "Normal Authentication", | |
| "Credential Stuffing", | |
| "Password Spray", | |
| "Brute Force", | |
| "Impossible Travel", | |
| "Insider Threat", | |
| "Privilege Abuse", | |
| "Suspicious Device" | |
| ]) | |
| self.features_used = config.get("classifier_features", [ | |
| "hour_of_day", "day_of_week", "is_weekend", "is_working_hour", | |
| "is_failure", "country_encoded", "is_mfa", "time_since_last_login", | |
| "rolling_failures_24h" | |
| ]) | |
| def train(self, X: pd.DataFrame, y: pd.Series): | |
| logger.info(f"Training Attack Classification model on {len(X)} samples...") | |
| missing_cols = [c for c in self.features_used if c not in X.columns] | |
| for c in missing_cols: | |
| X[c] = 0 | |
| X_train = X[self.features_used].fillna(0) | |
| self.model = xgb.XGBClassifier( | |
| objective="multi:softprob", | |
| num_class=len(self.labels), | |
| eval_metric="mlogloss", | |
| random_state=42 | |
| ) | |
| self.model.fit(X_train, y) | |
| self.persist() | |
| def persist(self): | |
| import hashlib | |
| self.model_dir.mkdir(parents=True, exist_ok=True) | |
| version_id = datetime.now().strftime("%Y%m%d_%H%M%S") | |
| model_path = self.model_dir / f"xgboost_attack_classifier_{version_id}.json" | |
| meta_path = self.model_dir / f"classifier_metadata_{version_id}.json" | |
| if self.model is not None: | |
| self.model.save_model(model_path) | |
| with open(model_path, "rb") as f: | |
| checksum = hashlib.sha256(f.read()).hexdigest() | |
| metadata = { | |
| "model_version": version_id, | |
| "training_dataset_version": "1.0", | |
| "feature_schema_version": "1.0", | |
| "training_date": datetime.now().isoformat(), | |
| "feature_schema": self.features_used, | |
| "labels": self.labels, | |
| "hyperparameters": self.model.get_params() if self.model else {}, | |
| "performance_metrics": {"mlogloss": "computed_during_eval"}, | |
| "model_checksum": checksum, | |
| "model_file": model_path.name | |
| } | |
| with open(meta_path, "w") as f: | |
| json.dump(metadata, f) | |
| self.version = version_id | |
| logger.info(f"Classifier model {version_id} and metadata persisted to {self.model_dir}") | |
| def load(self) -> bool: | |
| meta_files = sorted(self.model_dir.glob("classifier_metadata_*.json")) | |
| if not meta_files: | |
| return False | |
| latest_meta = meta_files[-1] | |
| with open(latest_meta, "r") as f: | |
| meta = json.load(f) | |
| self.labels = meta.get("labels", self.labels) | |
| self.version = meta.get("model_version", "1.0.0") | |
| model_file = meta.get("model_file") | |
| model_path = self.model_dir / model_file | |
| if model_path.exists(): | |
| self.model = xgb.XGBClassifier() | |
| self.model.load_model(model_path) | |
| logger.info(f"Pre-trained XGBoost classifier {self.version} loaded successfully.") | |
| return True | |
| return False | |
| def predict(self, features: pd.DataFrame) -> pd.DataFrame: | |
| if self.model is None: | |
| if not self.load(): | |
| raise RuntimeError("Classifier model is not loaded. Train or load first.") | |
| logger.info("Executing XGBoost Attack Classification inference...") | |
| # Prepare inference data | |
| missing_cols = [c for c in self.features_used if c not in features.columns] | |
| for c in missing_cols: | |
| features[c] = 0 | |
| X = features[self.features_used].fillna(0) | |
| probs = self.model.predict_proba(X) | |
| preds = probs.argmax(axis=1) | |
| confidences = probs.max(axis=1) | |
| predicted_labels = [self.labels[p] if p < len(self.labels) else "Unknown" for p in preds] | |
| return pd.DataFrame({ | |
| "event_id": features["event_id"], | |
| "attack_category": predicted_labels, | |
| "prediction_confidence": confidences, | |
| "classification_model_version": self.version | |
| }) | |