intelliscan/modules/classifier.py

"""
Module 4: ML Classifier (scan-result model, no data leakage)

Random Forest classifier trained on 11 behavioral features extracted from
live HTTP injection results. Avoids data leakage — no features derived
directly from the analyzer's signature outputs.

Behavioral features (11)
------------------------
Response features:
  1  response_len               length of response body (bytes)
  2  status_code                HTTP status code
  3  response_to_payload_ratio  len(response) / len(payload)
  4  has_html_tags              1 if response contains generic HTML markup
  5  response_length_anomaly    deviation from per-type mean response length

Payload features:
  6  payload_len                length of injected payload (chars)
  7  payload_special_chars_ratio  ratio of special chars in payload
  8  vuln_type_encoded          sqli=0, xss=1, lfi=2

URL features:
  9  url_length                 total length of target URL
  10 url_param_count            number of query parameters in URL
  11 url_depth                  directory depth of URL path
"""

from __future__ import annotations

import json
import logging
import statistics
from dataclasses import dataclass
from pathlib import Path
from urllib.parse import parse_qs, urlparse

import joblib
import numpy as np
from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import (
    accuracy_score,
    classification_report,
    confusion_matrix,
    f1_score,
    precision_score,
    recall_score,
)
from sklearn.model_selection import cross_val_score, train_test_split

from intelliscan.config import (
    CV_FOLDS,
    MODELS_DIR,
    RESULTS_DIR,
    RF_N_ESTIMATORS,
    RF_RANDOM_STATE,
    TRAIN_TEST_SPLIT,
)

log = logging.getLogger(__name__)

VULN_TYPE_MAP: dict[str, int] = {
    "sqli": 0,
    "xss": 1,
    "xss_r": 1,
    "xss_s": 1,
    "lfi": 2,
}

FEATURE_NAMES: tuple[str, ...] = (
    "response_len",
    "status_code",
    "response_to_payload_ratio",
    "has_html_tags",
    "response_length_anomaly",
    "payload_len",
    "payload_special_chars_ratio",
    "vuln_type_encoded",
    "url_length",
    "url_param_count",
    "url_depth",
)

SPECIAL_CHARS = set("'\"<>(){}[];%&|`$\\")


@dataclass
class TrainingReport:
    """Metrics produced after training."""

    accuracy: float
    precision: float
    recall: float
    f1_weighted: float
    cv_mean_f1: float
    cv_std_f1: float
    confusion_matrix: list[list[int]]
    classification_report: str
    feature_importances: dict[str, float]
    n_train: int
    n_test: int

    def summary(self) -> str:
        lines = [
            "================================================",
            "  Scan-Result ML Classifier — Training Report",
            "================================================",
            f"  Accuracy   : {self.accuracy  * 100:.1f}%",
            f"  Precision  : {self.precision * 100:.1f}%",
            f"  Recall     : {self.recall    * 100:.1f}%",
            f"  F1 (wgt)   : {self.f1_weighted * 100:.1f}%",
            f"  CV F1      : {self.cv_mean_f1:.2f}  +/-  {self.cv_std_f1:.2f}",
            f"  Train/Test : {self.n_train} / {self.n_test}",
            "------------------------------------------------",
            "  Feature importances (11 behavioral, no leakage):",
        ]
        for name, imp in sorted(self.feature_importances.items(), key=lambda x: -x[1])[:6]:
            lines.append(f"    {name:<32s} {imp:.3f}  ({imp*100:.1f}%)")
        lines.append("================================================")
        return "\n".join(lines)


class MLClassifier:
    """Random Forest wrapper using purely behavioral features."""

    def __init__(
        self,
        n_estimators: int = RF_N_ESTIMATORS,
        random_state: int = RF_RANDOM_STATE,
    ) -> None:
        self.model = RandomForestClassifier(
            n_estimators=n_estimators,
            random_state=random_state,
            class_weight="balanced",
            max_depth=10,  # prevent overfitting on small datasets
            min_samples_leaf=2,  # require at least 2 samples per leaf
            n_jobs=-1,
        )
        self.is_trained = False
        self._mean_response_len_per_type: dict[int, float] = {}

    # ── feature extraction ────────────────────────────────────────────────
    def _compute_per_type_means(self, dataset: list[dict]) -> None:
        """Compute mean response length per vuln_type (used for anomaly feature)."""
        per_type: dict[int, list[float]] = {0: [], 1: [], 2: []}
        for entry in dataset:
            vt = VULN_TYPE_MAP.get(entry.get("vuln_type", ""), 0)
            per_type[vt].append(float(entry.get("response_len", 0)))
        self._mean_response_len_per_type = {
            vt: statistics.mean(vals) if vals else 0.0 for vt, vals in per_type.items()
        }

    def extract_features(self, entry: dict) -> list[float]:
        response_len = float(entry.get("response_len", 0))
        payload = entry.get("payload", "") or ""
        payload_len = float(len(payload))
        body = entry.get("response_body") or ""
        vt = VULN_TYPE_MAP.get(entry.get("vuln_type", ""), 0)
        url = entry.get("target_url", "") or ""

        # F3: ratio response/payload (avoid division by zero)
        ratio = response_len / max(payload_len, 1.0)

        # F4: HTML markup presence (generic, not vuln-specific)
        body_lower = body.lower()
        has_html = (
            1.0
            if (
                "<html" in body_lower
                or "<body" in body_lower
                or "<div" in body_lower
                or "<p>" in body_lower
            )
            else 0.0
        )

        # F5: deviation from mean response length for this vuln type
        mean_len = self._mean_response_len_per_type.get(vt, response_len)
        anomaly = (response_len - mean_len) / max(mean_len, 1.0) if mean_len > 0 else 0.0

        # F7: ratio of special chars in payload
        special = sum(1 for c in payload if c in SPECIAL_CHARS)
        special_ratio = special / max(payload_len, 1.0)

        # F9–F11: URL-based features
        url_length = float(len(url))
        try:
            parsed = urlparse(url)
            query_params = parse_qs(parsed.query, keep_blank_values=True)
            url_param_count = float(len(query_params))
            url_depth = float(max(0, len([s for s in parsed.path.split("/") if s])))
        except Exception:
            url_param_count = 0.0
            url_depth = 0.0

        return [
            response_len,  # F1
            float(entry.get("status_code", 200)),  # F2
            float(ratio),  # F3
            has_html,  # F4
            float(anomaly),  # F5
            payload_len,  # F6
            float(special_ratio),  # F7
            float(vt),  # F8
            url_length,  # F9
            url_param_count,  # F10
            url_depth,  # F11
        ]

    # ── training ──────────────────────────────────────────────────────────
    def train(self, dataset: list[dict]) -> TrainingReport:
        if len(dataset) < 4:
            raise ValueError(f"Need at least 4 samples, got {len(dataset)}")

        y = np.array([1 if e["label"] == "VULNERABLE" else 0 for e in dataset])

        # Split BEFORE computing per-type response-length means so the anomaly
        # feature (F5) never sees test-set statistics (no train/test leakage).
        indices = np.arange(len(dataset))
        train_idx, test_idx = train_test_split(
            indices,
            test_size=TRAIN_TEST_SPLIT,
            random_state=RF_RANDOM_STATE,
            stratify=y,
        )
        self._compute_per_type_means([dataset[i] for i in train_idx])

        X = np.array([self.extract_features(e) for e in dataset])
        X_train, X_test = X[train_idx], X[test_idx]
        y_train, y_test = y[train_idx], y[test_idx]

        log.info("Training RF on %d samples (11 behavioral features) ...", len(X_train))
        self.model.fit(X_train, y_train)
        self.is_trained = True

        y_pred = self.model.predict(X_test)
        cv_scores = cross_val_score(self.model, X, y, cv=CV_FOLDS, scoring="f1_weighted")

        importances = dict(zip(FEATURE_NAMES, self.model.feature_importances_, strict=True))

        report = TrainingReport(
            accuracy=accuracy_score(y_test, y_pred),
            precision=precision_score(y_test, y_pred, zero_division=0),
            recall=recall_score(y_test, y_pred, zero_division=0),
            f1_weighted=f1_score(y_test, y_pred, average="weighted", zero_division=0),
            cv_mean_f1=float(cv_scores.mean()),
            cv_std_f1=float(cv_scores.std()),
            confusion_matrix=confusion_matrix(y_test, y_pred).tolist(),
            classification_report=classification_report(
                y_test,
                y_pred,
                target_names=["NOT_VULNERABLE", "VULNERABLE"],
                zero_division=0,
            ),
            feature_importances=importances,
            n_train=len(X_train),
            n_test=len(X_test),
        )
        log.info("\n%s", report.summary())
        return report

    # ── inference ─────────────────────────────────────────────────────────
    def predict(self, entry: dict) -> tuple[str, float]:
        if not self.is_trained:
            raise RuntimeError("Model not trained. Call train() or load() first.")
        X = np.array([self.extract_features(entry)])
        pred = int(self.model.predict(X)[0])
        proba = float(self.model.predict_proba(X)[0][pred])
        return ("VULNERABLE" if pred == 1 else "NOT_VULNERABLE"), proba

    # ── persistence ───────────────────────────────────────────────────────
    def save(self, path: str | Path = MODELS_DIR / "model.pkl") -> Path:
        path = Path(path)
        path.parent.mkdir(parents=True, exist_ok=True)
        joblib.dump(
            {
                "model": self.model,
                "mean_response_len_per_type": self._mean_response_len_per_type,
            },
            path,
        )
        log.info("Model saved -> %s", path)
        return path

    def load(self, path: str | Path = MODELS_DIR / "model.pkl") -> None:
        data = joblib.load(Path(path))
        if isinstance(data, dict):
            self.model = data["model"]
            self._mean_response_len_per_type = data.get("mean_response_len_per_type", {})
        else:
            # Backward compatibility with old model files
            self.model = data
        self.is_trained = True
        log.info("Model loaded <- %s", path)


# ── CLI helper ────────────────────────────────────────────────────────────
def main(input_file: str = "labeled_results.json") -> None:
    path = RESULTS_DIR / input_file
    dataset = json.loads(path.read_text())
    clf = MLClassifier()
    report = clf.train(dataset)
    clf.save()
    print(report.summary())


if __name__ == "__main__":
    import argparse

    p = argparse.ArgumentParser(description="IntelliScan ML classifier")
    p.add_argument("--input", default="labeled_results.json")
    args = p.parse_args()

    logging.basicConfig(level=logging.INFO, format="%(levelname)s | %(message)s")
    main(args.input)