Spaces:
Paused
Paused
File size: 4,891 Bytes
5a3b9db | 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 | 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
})
|