File size: 4,041 Bytes
8e95c1c | 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 | 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)
|