diff --git a/src/codeguard/core/models.py b/src/codeguard/core/models.py new file mode 100644 index 0000000..112b3e7 --- /dev/null +++ b/src/codeguard/core/models.py @@ -0,0 +1,72 @@ +"""Data models for CodeGuard.""" + +from enum import Enum +from typing import Optional +from pydantic import BaseModel, Field + + +class Severity(str, Enum): + CRITICAL = "critical" + HIGH = "high" + MEDIUM = "medium" + LOW = "low" + INFO = "info" + + +class Language(str, Enum): + PYTHON = "python" + JAVASCRIPT = "javascript" + TYPESCRIPT = "typescript" + GO = "go" + RUST = "rust" + + +class FindingType(str, Enum): + VULNERABILITY = "vulnerability" + CODE_SMELL = "code_smell" + BUG = "bug" + SECURITY_PATTERN = "security_pattern" + + +class Location(BaseModel): + file: str = Field(..., description="File path") + line: int = Field(..., description="Line number") + end_line: Optional[int] = Field(None, description="End line number") + column: Optional[int] = Field(None, description="Column number") + code_snippet: Optional[str] = Field(None, description="Code snippet at location") + + +class FixSuggestion(BaseModel): + description: str = Field(..., description="Description of the fix") + code: Optional[str] = Field(None, description="Suggested code replacement") + explanation: Optional[str] = Field(None, json_schema_extra={"description": "Why this fix works"}) + + +class Finding(BaseModel): + id: str = Field(..., description="Unique identifier") + type: FindingType = Field(..., description="Type of finding") + severity: Severity = Field(..., description="Severity level") + title: str = Field(..., description="Brief title") + description: str = Field(..., description="Detailed description") + location: Location = Field(..., description="Location in code") + cwe_id: Optional[str] = Field(None, description="CWE ID if applicable") + fix: Optional[FixSuggestion] = Field(None, description="Suggested fix") + language: Optional[Language] = Field(None, description="Programming language") + + +class ScanResult(BaseModel): + files_scanned: int = Field(..., description="Number of files scanned") + findings: list[Finding] = Field(default_factory=list, description="Findings") + duration_seconds: float = Field(..., description="Scan duration") + language_breakdown: dict[str, int] = Field(default_factory=dict) + severity_breakdown: dict[str, int] = Field(default_factory=dict) + + +class Config(BaseModel): + model: str = Field(default="codellama", description="Ollama model") + languages: list[Language] = Field(default_factory=lambda: [Language.PYTHON]) + severity_threshold: Severity = Field(default=Severity.INFO, description="Minimum severity to report") + fail_on_critical: bool = Field(default=False, description="Fail on critical findings") + max_file_size: int = Field(default=100000, description="Max file size in bytes") + chunk_size: int = Field(default=8000, description="Chunk size for large files") + custom_rules: list[str] = Field(default_factory=list, description="Custom rule paths")