Spaces:
Runtime error
Runtime error
| """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": "<html><body>data</body></html>", | |
| "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": "<html><body>test</body></html>", | |
| } | |
| 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": "<html>data</html>", | |
| "label": "VULNERABLE", | |
| } | |
| ) | |
| dataset.append( | |
| { | |
| "vuln_type": "sqli", | |
| "payload": "x", | |
| "status_code": 200, | |
| "response_len": 1000, | |
| "response_body": "<html>no</html>", | |
| "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) | |