CosmoGrepperAI/model.py at main · appsecninja32/CosmoGrepperAI · GitHub

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
from pydantic import BaseModel
from typing import Optional


# ---------------------------------------------------------
# ENRICHMENT MODELS
# ---------------------------------------------------------

class CVSS(BaseModel):
    base_score: float = 0.0
    vector: str = ""


class OWASP(BaseModel):
    top10_category: str = ""


class DataClassification(BaseModel):
    level: str = ""


class Risk(BaseModel):
    risk_score: float = 0.0
    likelihood: float = 0.0
    impact: float = 0.0


# ---------------------------------------------------------
# FINDING MODEL
# ---------------------------------------------------------

class Finding(BaseModel):
    # Required fields
    source: str
    file: str
    line: int
    title: str
    description: str

    # Optional fields (fixes missing "type")
    type: Optional[str] = "semgrep"

    # Enrichment fields
    severity_score: int = 0
    cvss: CVSS = CVSS()
    owasp: OWASP = OWASP()
    data_classification: DataClassification = DataClassification()
    risk: Risk = Risk()

    # Snippet support
    snippet: Optional[str] = None

    # LLM FP Analysis
    false_positive: bool = False
    fp_reason: str = ""
    vulnerability_explanation: str = ""
    mitigation: str = ""


# ---------------------------------------------------------
# UNIVERSAL FACTORY FUNCTION
# ---------------------------------------------------------

def new_finding(**kwargs):
    """
    Universal factory for creating Finding objects.
    Accepts ANY keyword arguments so the scan engine
    can pass severity_score, snippet, source, etc.
    """

    # Convert severity strings → numeric scores
    sev = kwargs.get("severity_score")

    if isinstance(sev, str):
        sev_map = {
            "ERROR": 5,
            "WARNING": 3,
            "INFO": 1
        }
        kwargs["severity_score"] = sev_map.get(sev.upper(), 1)

    return Finding(**kwargs)