"""Tests for the ML Classifier module (behavioral features only).""" import pytest from intelliscan.modules.classifier import ( FEATURE_NAMES, VULN_TYPE_MAP, MLClassifier, ) def test_feature_extraction_returns_11_floats(): clf = MLClassifier() clf._mean_response_len_per_type = {0: 4500.0, 1: 3800.0, 2: 5000.0} entry = { "vuln_type": "sqli", "payload": "' OR 1=1-- ", "status_code": 200, "response_len": 4821, "response_body": "data", "target_url": "http://localhost:8080/vulnerabilities/sqli/?id=1", } features = clf.extract_features(entry) assert len(features) == len(FEATURE_NAMES) == 11 assert all(isinstance(f, float) for f in features) def test_html_detection(): clf = MLClassifier() clf._mean_response_len_per_type = {0: 1000.0, 1: 1000.0, 2: 1000.0} entry_with_html = { "vuln_type": "sqli", "payload": "x", "status_code": 200, "response_len": 1000, "response_body": "test", } entry_without_html = { "vuln_type": "sqli", "payload": "x", "status_code": 200, "response_len": 1000, "response_body": "plain text response", } has_html_idx = FEATURE_NAMES.index("has_html_tags") assert clf.extract_features(entry_with_html)[has_html_idx] == 1.0 assert clf.extract_features(entry_without_html)[has_html_idx] == 0.0 def test_vuln_type_encoding(): assert VULN_TYPE_MAP["sqli"] == 0 assert VULN_TYPE_MAP["xss"] == 1 assert VULN_TYPE_MAP["xss_r"] == 1 assert VULN_TYPE_MAP["xss_s"] == 1 assert VULN_TYPE_MAP["lfi"] == 2 def test_train_too_small_dataset_raises(): clf = MLClassifier() with pytest.raises(ValueError): clf.train([{"label": "VULNERABLE"}]) def test_predict_before_training_raises(): clf = MLClassifier() with pytest.raises(RuntimeError): clf.predict({"vuln_type": "sqli", "payload": "x", "response_body": ""}) def test_train_with_synthetic_data(): """End-to-end training on a tiny synthetic dataset.""" dataset = [] for _ in range(8): dataset.append( { "vuln_type": "sqli", "payload": "' OR 1=1-- ", "status_code": 200, "response_len": 5000, "response_body": "data", "label": "VULNERABLE", } ) dataset.append( { "vuln_type": "sqli", "payload": "x", "status_code": 200, "response_len": 1000, "response_body": "no", "label": "NOT_VULNERABLE", } ) clf = MLClassifier() report = clf.train(dataset) assert 0 <= report.accuracy <= 1 assert 0 <= report.f1_weighted <= 1 assert sum(report.feature_importances.values()) == pytest.approx(1.0, abs=0.01)