| import argparse |
| import sys |
| from pathlib import Path |
|
|
| import joblib |
| import numpy as np |
| from sklearn.ensemble import RandomForestClassifier |
| from sklearn.model_selection import StratifiedKFold, cross_validate |
| from sklearn.metrics import classification_report, roc_auc_score |
|
|
| sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) |
|
|
| from app.engine.ml_classifier import extract_features, FEATURE_NAMES |
|
|
|
|
| def load_samples(data_dir: Path) -> tuple[np.ndarray, np.ndarray]: |
| benign_dir = data_dir / "benign" |
| malicious_dir = data_dir / "malicious" |
|
|
| X_list = [] |
| y_list = [] |
| skipped = 0 |
|
|
| for label, directory in [(0, benign_dir), (1, malicious_dir)]: |
| if not directory.exists(): |
| print(f"Directory not found: {directory}") |
| continue |
| for filepath in directory.iterdir(): |
| if not filepath.is_file(): |
| continue |
| try: |
| file_bytes = filepath.read_bytes() |
| features = extract_features(file_bytes) |
| if features is not None: |
| X_list.append(features) |
| y_list.append(label) |
| else: |
| skipped += 1 |
| except Exception: |
| skipped += 1 |
|
|
| print(f"Loaded {len(X_list)} samples ({skipped} skipped)") |
| return np.array(X_list), np.array(y_list) |
|
|
|
|
| def train(data_dir: str, output_path: str, model_type: str = "rf") -> None: |
| X, y = load_samples(Path(data_dir)) |
|
|
| if len(X) < 20: |
| print("Not enough samples for training. Need at least 20.") |
| sys.exit(1) |
|
|
| benign_count = int(np.sum(y == 0)) |
| malicious_count = int(np.sum(y == 1)) |
| print(f"Benign: {benign_count}, Malicious: {malicious_count}") |
|
|
| if model_type == "catboost": |
| try: |
| from catboost import CatBoostClassifier |
| model = CatBoostClassifier( |
| iterations=500, |
| learning_rate=0.05, |
| depth=6, |
| l2_leaf_reg=5, |
| border_count=128, |
| class_weights={0: 1.0, 1: 2.0}, |
| eval_metric="Precision", |
| random_seed=42, |
| verbose=0, |
| ) |
| except ImportError: |
| print("CatBoost not installed, falling back to RandomForest") |
| model_type = "rf" |
|
|
| if model_type == "rf": |
| model = RandomForestClassifier( |
| n_estimators=100, |
| max_depth=10, |
| random_state=42, |
| ) |
|
|
| cv = StratifiedKFold(n_splits=5, shuffle=True, random_state=42) |
| cv_results = cross_validate( |
| model, X, y, cv=cv, |
| scoring=["precision", "recall", "f1", "roc_auc"], |
| return_train_score=False, |
| ) |
|
|
| print("\n5-Fold Cross-Validation Results:") |
| print(f" Precision: {cv_results['test_precision'].mean():.4f} (+/- {cv_results['test_precision'].std():.4f})") |
| print(f" Recall: {cv_results['test_recall'].mean():.4f} (+/- {cv_results['test_recall'].std():.4f})") |
| print(f" F1: {cv_results['test_f1'].mean():.4f} (+/- {cv_results['test_f1'].std():.4f})") |
| print(f" AUC-ROC: {cv_results['test_roc_auc'].mean():.4f} (+/- {cv_results['test_roc_auc'].std():.4f})") |
|
|
| model.fit(X, y) |
|
|
| output = Path(output_path) |
| output.parent.mkdir(parents=True, exist_ok=True) |
| joblib.dump(model, output) |
| print(f"\nModel saved to {output}") |
|
|
| if hasattr(model, "feature_importances_"): |
| print("\nFeature Importances:") |
| importances = model.feature_importances_ |
| sorted_idx = np.argsort(importances)[::-1] |
| for idx in sorted_idx: |
| print(f" {FEATURE_NAMES[idx]:30s} {importances[idx]:.4f}") |
|
|
|
|
| if __name__ == "__main__": |
| parser = argparse.ArgumentParser() |
| parser.add_argument("--data-dir", required=True) |
| parser.add_argument("--output", default="app/models/rf_model.pkl") |
| parser.add_argument("--model-type", choices=["rf", "catboost"], default="rf") |
| args = parser.parse_args() |
| train(args.data_dir, args.output, args.model_type) |
|
|
