From fa143e76333c493515c64e9b4d788328fde302ee Mon Sep 17 00:00:00 2001 From: fatihbugrakdogan Date: Fri, 6 Mar 2026 22:25:47 +0300 Subject: [PATCH 1/2] cvss implementation for severity added --- src/kai/definitions/exploit/parsers.py | 26 ++++- src/kai/definitions/exploit/prompt.py | 86 ++++++++++++--- src/kai/state/models.py | 12 +++ src/kai/utils/cvss.py | 141 +++++++++++++++++++++++++ tests/test_cvss.py | 116 ++++++++++++++++++++ 5 files changed, 363 insertions(+), 18 deletions(-) create mode 100644 src/kai/utils/cvss.py create mode 100644 tests/test_cvss.py diff --git a/src/kai/definitions/exploit/parsers.py b/src/kai/definitions/exploit/parsers.py index e56dff5c..b44f8d4c 100644 --- a/src/kai/definitions/exploit/parsers.py +++ b/src/kai/definitions/exploit/parsers.py @@ -19,6 +19,7 @@ from kai import generate_id from kai.state.base import StateManager from kai.state.models import ExploitRecord, FixRecord +from kai.utils.cvss import validate_and_compute log = logging.getLogger(__name__) @@ -113,11 +114,30 @@ def process_fixer_result( return raw_result exploit_id = existing.exploit_id + # Validate CVSS vector and derive score + severity + cvss = validate_and_compute(result.get("cvss_vector")) + if cvss is not None: + severity = cvss.severity + cvss_vector = cvss.vector + cvss_score = cvss.score + else: + severity = result.get("severity", "") + cvss_vector = result.get("cvss_vector") + cvss_score = None + if cvss_vector: + log.warning( + "Invalid CVSS vector from fixer: %s — " + "falling back to LLM-provided severity", + cvss_vector, + ) + state_manager.update_exploit( run_id, exploit_id, status="verified_and_fixed", - severity=result.get("severity"), + severity=severity, + cvss_vector=cvss_vector, + cvss_score=cvss_score, patch=result.get("patch"), test_results=result.get("test_results"), ) @@ -130,7 +150,9 @@ def process_fixer_result( hypothesis=result.get("hypothesis", ""), file=result.get("file", ""), function=result.get("function", ""), - severity=result.get("severity", ""), + severity=severity or "", + cvss_vector=cvss_vector or "", + cvss_score=cvss_score, patch=result.get("patch", ""), test_results=result.get("test_results", ""), applied=False, diff --git a/src/kai/definitions/exploit/prompt.py b/src/kai/definitions/exploit/prompt.py index 32941134..d52fbdf6 100644 --- a/src/kai/definitions/exploit/prompt.py +++ b/src/kai/definitions/exploit/prompt.py @@ -54,8 +54,9 @@ function: str, poc_code: str, test_output: str, \ affected_files: list[str] = []) -> str` Generate a minimal patch for a verified exploit and assess \ - its severity. Returns a dict with: severity (Critical / \ - High / Medium / Low), patch, and test results. + its severity via CVSS v3.1. Returns a dict with: severity \ + (Critical / High / Medium / Low), cvss_vector, patch, and \ + test results. - `spawn_researcher(query: str) -> str` Research vulnerability patterns relevant to the target \ codebase. When given a structural summary (project type, \ @@ -250,6 +251,7 @@ "file": str, "function": str, "severity": str, # assigned by fixer + "cvss_vector": str, # CVSS:3.1/... from fixer "poc_code": str, # from verifier "patch": str, # from fixer "test_results": str, # from fixer @@ -820,19 +822,70 @@ Write a minimal patch that neutralizes the exploit, verify the \ patch by running tests (the PoC should now fail, existing tests \ - should still pass), and assess severity. Hashline editing tools \ - use refs from `read_file` output — re-read after edits for fresh \ - refs. - - Severity scale: - - **Critical**: direct loss of funds, total access-control bypass, \ - or protocol-breaking invariant violation - - **High**: conditional fund loss, privilege escalation, or state \ - corruption exploitable under realistic conditions - - **Medium**: limited impact requiring specific preconditions, \ - DoS, or griefing attacks - - **Low**: informational, gas inefficiency, or edge-case-only \ - issues unlikely to be exploited + should still pass), and assess severity using CVSS v3.1. \ + Hashline editing tools use refs from `read_file` output — \ + re-read after edits for fresh refs. + + ## CVSS v3.1 scoring + + You MUST produce a CVSS v3.1 vector string by evaluating these \ + eight base metrics for the vulnerability: + + **Attack Vector (AV)** — how the attacker reaches the target: + - N (Network): exploitable over the network (e.g. HTTP request) + - A (Adjacent): requires same LAN/broadcast domain + - L (Local): requires local access (e.g. CLI, local file) + - P (Physical): requires physical hardware access + + **Attack Complexity (AC)** — conditions beyond attacker control: + - L (Low): no special conditions, reproducible at will + - H (High): requires specific configuration, race condition, \ + or non-default setup + + **Privileges Required (PR)** — auth level needed: + - N (None): unauthenticated / any external caller + - L (Low): basic user-level privileges + - H (High): admin/owner privileges required + + **User Interaction (UI)** — does a victim need to act? + - N (None): no user interaction needed + - R (Required): victim must click/approve/interact + + **Scope (S)** — does the exploit cross a trust boundary? + - U (Unchanged): impact stays within the vulnerable component + - C (Changed): impacts resources beyond the vulnerable component + + **Confidentiality (C)** — impact on data secrecy: + - N (None): no confidentiality loss + - L (Low): limited data exposure + - H (High): total confidentiality loss / all data exposed + + **Integrity (I)** — impact on data correctness: + - N (None): no integrity impact + - L (Low): limited data modification possible + - H (High): total integrity loss / arbitrary data modification + + **Availability (A)** — impact on system uptime: + - N (None): no availability impact + - L (Low): degraded performance / partial disruption + - H (High): total denial of service / system shutdown + + Assemble the vector as: \ + `CVSS:3.1/AV:X/AC:X/PR:X/UI:X/S:X/C:X/I:X/A:X` + + The numeric score is computed from the vector automatically — \ + you do NOT need to calculate it yourself. Just provide the \ + vector string. + + Severity maps from the computed score: + - **Critical** (9.0–10.0): direct loss of funds, total \ + access-control bypass, protocol-breaking invariant violation + - **High** (7.0–8.9): conditional fund loss, privilege \ + escalation, state corruption under realistic conditions + - **Medium** (4.0–6.9): limited impact with specific \ + preconditions, DoS, griefing attacks + - **Low** (0.1–3.9): informational, gas inefficiency, \ + edge-case-only issues unlikely to be exploited ## Capturing the patch @@ -855,7 +908,8 @@ "hypothesis": context["hypothesis"], "file": context["file"], "function": context["function"], - "severity": "Critical", # your assessment + "severity": "Critical", # derived from cvss_score + "cvss_vector": "CVSS:3.1/AV:N/AC:L/PR:N/UI:N/S:U/C:H/I:H/A:H", "poc_code": context["poc_code"], "patch": patch_text, # unified diff from git diff "test_results": "...", # output from running tests after fix diff --git a/src/kai/state/models.py b/src/kai/state/models.py index 438d9db7..ac2e9b48 100644 --- a/src/kai/state/models.py +++ b/src/kai/state/models.py @@ -117,6 +117,8 @@ class ExploitRecord: poc_code: str | None = None test_output: str | None = None severity: str | None = None + cvss_vector: str | None = None + cvss_score: float | None = None patch: str | None = None test_results: str | None = None @@ -136,6 +138,8 @@ def to_dict(self) -> dict[str, Any]: "poc_code": self.poc_code, "test_output": self.test_output, "severity": self.severity, + "cvss_vector": self.cvss_vector, + "cvss_score": self.cvss_score, "patch": self.patch, "test_results": self.test_results, } @@ -157,6 +161,8 @@ def from_dict(cls, data: dict[str, Any]) -> ExploitRecord: poc_code=data.get("poc_code"), test_output=data.get("test_output"), severity=data.get("severity"), + cvss_vector=data.get("cvss_vector"), + cvss_score=data.get("cvss_score"), patch=data.get("patch"), test_results=data.get("test_results"), ) @@ -176,6 +182,8 @@ class FixRecord: severity: str patch: str test_results: str + cvss_vector: str = "" + cvss_score: float | None = None applied: bool = False def to_dict(self) -> dict[str, Any]: @@ -189,6 +197,8 @@ def to_dict(self) -> dict[str, Any]: "file": self.file, "function": self.function, "severity": self.severity, + "cvss_vector": self.cvss_vector, + "cvss_score": self.cvss_score, "patch": self.patch, "test_results": self.test_results, "applied": self.applied, @@ -206,6 +216,8 @@ def from_dict(cls, data: dict[str, Any]) -> FixRecord: file=data["file"], function=data["function"], severity=data["severity"], + cvss_vector=data.get("cvss_vector", ""), + cvss_score=data.get("cvss_score"), patch=data["patch"], test_results=data["test_results"], applied=data.get("applied", False), diff --git a/src/kai/utils/cvss.py b/src/kai/utils/cvss.py new file mode 100644 index 00000000..92a5f090 --- /dev/null +++ b/src/kai/utils/cvss.py @@ -0,0 +1,141 @@ +"""CVSS v3.1 Base Score calculator and vector string parser. + +Pure-Python implementation — no external dependencies. +Spec: https://www.first.org/cvss/v3.1/specification-document +""" + +from __future__ import annotations + +import math +import re +from dataclasses import dataclass + +_VECTOR_RE = re.compile( + r"^CVSS:3\.1" + r"/AV:([NALP])" + r"/AC:([LH])" + r"/PR:([NLH])" + r"/UI:([NR])" + r"/S:([UC])" + r"/C:([NLH])" + r"/I:([NLH])" + r"/A:([NLH])$" +) + +# Metric value → numeric weight (CVSS v3.1 spec tables) +_AV = {"N": 0.85, "A": 0.62, "L": 0.55, "P": 0.20} +_AC = {"L": 0.77, "H": 0.44} +_UI = {"N": 0.85, "R": 0.62} + +# PR depends on Scope +_PR_UNCHANGED = {"N": 0.85, "L": 0.62, "H": 0.27} +_PR_CHANGED = {"N": 0.85, "L": 0.68, "H": 0.50} + +_CIA = {"N": 0.0, "L": 0.22, "H": 0.56} + +SEVERITY_RANGES = { + "Critical": (9.0, 10.0), + "High": (7.0, 8.9), + "Medium": (4.0, 6.9), + "Low": (0.1, 3.9), + "None": (0.0, 0.0), +} + + +def _roundup(x: float) -> float: + """CVSS v3.1 roundup: smallest tenth >= x.""" + return math.ceil(x * 10) / 10 + + +@dataclass(frozen=True) +class CVSSResult: + vector: str + score: float + severity: str + + +def parse_vector(vector: str) -> dict[str, str] | None: + """Parse a CVSS:3.1 vector string into metric components. + + Returns None if the vector is malformed. + """ + m = _VECTOR_RE.match(vector.strip()) + if m is None: + return None + return { + "AV": m.group(1), + "AC": m.group(2), + "PR": m.group(3), + "UI": m.group(4), + "S": m.group(5), + "C": m.group(6), + "I": m.group(7), + "A": m.group(8), + } + + +def compute_score(metrics: dict[str, str]) -> float: + """Compute CVSS v3.1 base score from metric dict.""" + scope_changed = metrics["S"] == "C" + pr_table = _PR_CHANGED if scope_changed else _PR_UNCHANGED + + isc_base = 1.0 - ( + (1.0 - _CIA[metrics["C"]]) + * (1.0 - _CIA[metrics["I"]]) + * (1.0 - _CIA[metrics["A"]]) + ) + + if scope_changed: + impact = 7.52 * (isc_base - 0.029) - 3.25 * (isc_base - 0.02) ** 15 + else: + impact = 6.42 * isc_base + + if impact <= 0: + return 0.0 + + exploitability = ( + 8.22 + * _AV[metrics["AV"]] + * _AC[metrics["AC"]] + * pr_table[metrics["PR"]] + * _UI[metrics["UI"]] + ) + + if scope_changed: + score = _roundup(min(1.08 * (impact + exploitability), 10.0)) + else: + score = _roundup(min(impact + exploitability, 10.0)) + + return score + + +def score_to_severity(score: float) -> str: + """Map a numeric CVSS score to a severity label.""" + if score >= 9.0: + return "Critical" + if score >= 7.0: + return "High" + if score >= 4.0: + return "Medium" + if score > 0.0: + return "Low" + return "None" + + +def validate_and_compute( + vector: str | None, + claimed_score: float | None = None, +) -> CVSSResult | None: + """Validate a CVSS vector and compute/correct the score. + + Returns a CVSSResult with the correct score derived from the + vector, or None if the vector is invalid/missing. + """ + if not vector: + return None + metrics = parse_vector(vector) + if metrics is None: + return None + score = compute_score(metrics) + severity = score_to_severity(score) + return CVSSResult(vector=vector.strip(), score=score, severity=severity) diff --git a/tests/test_cvss.py b/tests/test_cvss.py new file mode 100644 index 00000000..cd74221e --- /dev/null +++ b/tests/test_cvss.py @@ -0,0 +1,116 @@ +"""Tests for the CVSS v3.1 calculator.""" + +import pytest + +from kai.utils.cvss import ( + CVSSResult, + compute_score, + parse_vector, + score_to_severity, + validate_and_compute, +) + + +class TestParseVector: + def test_valid_vector(self): + result = parse_vector("CVSS:3.1/AV:N/AC:L/PR:N/UI:N/S:U/C:H/I:H/A:H") + assert result == { + "AV": "N", "AC": "L", "PR": "N", "UI": "N", + "S": "U", "C": "H", "I": "H", "A": "H", + } + + def test_invalid_prefix(self): + assert parse_vector("CVSS:2.0/AV:N/AC:L/PR:N/UI:N/S:U/C:H/I:H/A:H") is None + + def test_missing_metric(self): + assert parse_vector("CVSS:3.1/AV:N/AC:L/PR:N/UI:N/S:U/C:H/I:H") is None + + def test_invalid_metric_value(self): + assert parse_vector("CVSS:3.1/AV:X/AC:L/PR:N/UI:N/S:U/C:H/I:H/A:H") is None + + def test_empty_string(self): + assert parse_vector("") is None + + def test_strips_whitespace(self): + result = parse_vector(" CVSS:3.1/AV:N/AC:L/PR:N/UI:N/S:U/C:H/I:H/A:H ") + assert result is not None + + +class TestComputeScore: + """Test against known CVSS v3.1 reference scores.""" + + def test_max_score(self): + # AV:N/AC:L/PR:N/UI:N/S:C/C:H/I:H/A:H = 10.0 + metrics = {"AV": "N", "AC": "L", "PR": "N", "UI": "N", + "S": "C", "C": "H", "I": "H", "A": "H"} + assert compute_score(metrics) == 10.0 + + def test_scope_unchanged_high(self): + # AV:N/AC:L/PR:N/UI:N/S:U/C:H/I:H/A:H = 9.8 + metrics = {"AV": "N", "AC": "L", "PR": "N", "UI": "N", + "S": "U", "C": "H", "I": "H", "A": "H"} + assert compute_score(metrics) == 9.8 + + def test_zero_impact(self): + # All CIA = None → impact = 0 → score = 0 + metrics = {"AV": "N", "AC": "L", "PR": "N", "UI": "N", + "S": "U", "C": "N", "I": "N", "A": "N"} + assert compute_score(metrics) == 0.0 + + def test_medium_score(self): + # AV:N/AC:H/PR:L/UI:R/S:U/C:L/I:L/A:N = 3.7 + metrics = {"AV": "N", "AC": "H", "PR": "L", "UI": "R", + "S": "U", "C": "L", "I": "L", "A": "N"} + score = compute_score(metrics) + assert 3.0 <= score <= 4.0 + + def test_scope_changed_medium(self): + # AV:N/AC:L/PR:N/UI:R/S:C/C:L/I:L/A:N = 6.1 + metrics = {"AV": "N", "AC": "L", "PR": "N", "UI": "R", + "S": "C", "C": "L", "I": "L", "A": "N"} + assert compute_score(metrics) == 6.1 + + +class TestScoreToSeverity: + @pytest.mark.parametrize("score,expected", [ + (0.0, "None"), + (0.1, "Low"), + (3.9, "Low"), + (4.0, "Medium"), + (6.9, "Medium"), + (7.0, "High"), + (8.9, "High"), + (9.0, "Critical"), + (10.0, "Critical"), + ]) + def test_severity_ranges(self, score, expected): + assert score_to_severity(score) == expected + + +class TestValidateAndCompute: + def test_valid_vector(self): + result = validate_and_compute( + "CVSS:3.1/AV:N/AC:L/PR:N/UI:N/S:U/C:H/I:H/A:H" + ) + assert result is not None + assert result.score == 9.8 + assert result.severity == "Critical" + assert result.vector == "CVSS:3.1/AV:N/AC:L/PR:N/UI:N/S:U/C:H/I:H/A:H" + + def test_none_vector(self): + assert validate_and_compute(None) is None + + def test_empty_vector(self): + assert validate_and_compute("") is None + + def test_invalid_vector(self): + assert validate_and_compute("not-a-vector") is None + + def test_ignores_claimed_score(self): + # Even if claimed score is wrong, computed score is correct + result = validate_and_compute( + "CVSS:3.1/AV:N/AC:L/PR:N/UI:N/S:U/C:H/I:H/A:H", + claimed_score=5.0, + ) + assert result is not None + assert result.score == 9.8 From c92f342c9697e5b8b0f005ff452bfbf4b1fa0122 Mon Sep 17 00:00:00 2001 From: fatihbugrakdogan Date: Sat, 7 Mar 2026 00:43:11 +0300 Subject: [PATCH 2/2] ruff and lint check --- tests/test_cvss.py | 1 - tests/test_parsing.py | 1 - 2 files changed, 2 deletions(-) diff --git a/tests/test_cvss.py b/tests/test_cvss.py index cd74221e..279ce892 100644 --- a/tests/test_cvss.py +++ b/tests/test_cvss.py @@ -3,7 +3,6 @@ import pytest from kai.utils.cvss import ( - CVSSResult, compute_score, parse_vector, score_to_severity, diff --git a/tests/test_parsing.py b/tests/test_parsing.py index 53ebf89c..08b826d5 100644 --- a/tests/test_parsing.py +++ b/tests/test_parsing.py @@ -4,7 +4,6 @@ import json -import pytest from ra.utils.parsing import find_code_blocks, find_final_answer