diff --git a/docksec/cli.py b/docksec/cli.py index ccdb8c3..f3b9993 100644 --- a/docksec/cli.py +++ b/docksec/cli.py @@ -3,6 +3,7 @@ import sys import os import argparse +from docksec.enums import LLMProvider def get_version() -> str: """Return the installed package version. @@ -42,7 +43,7 @@ def main() -> None: parser.add_argument('--ai-only', action='store_true', help='Run only AI-based recommendations (requires Dockerfile)') parser.add_argument('--scan-only', action='store_true', help='Run only Dockerfile/image scanning (requires --image)') parser.add_argument('--image-only', action='store_true', help='Scan only the Docker image without Dockerfile analysis') - parser.add_argument('--provider', choices=['openai', 'anthropic', 'google', 'ollama'], + parser.add_argument('--provider', choices=LLMProvider.values(), help='LLM provider to use (default: openai, can also set LLM_PROVIDER env var)') parser.add_argument('--model', help='Model name to use (e.g., gpt-4o, claude-3-5-sonnet-20241022, gemini-1.5-pro, llama3.1)') parser.add_argument('--version', action='version', version=f'DockSec {get_version()}') diff --git a/docksec/config_manager.py b/docksec/config_manager.py index 9bb6cf5..48df4a2 100644 --- a/docksec/config_manager.py +++ b/docksec/config_manager.py @@ -10,6 +10,7 @@ import logging from typing import Optional, Any, Dict from dataclasses import dataclass, field +from docksec.enums import LLMProvider, Severity # Set up logger logger = logging.getLogger(__name__) @@ -119,7 +120,7 @@ def _validate(self) -> None: ValueError: If configuration values are invalid """ # Validate LLM provider - valid_providers = ['openai', 'anthropic', 'google', 'ollama'] + valid_providers = LLMProvider.values() if self.llm_provider not in valid_providers: raise ValueError(f"Invalid llm_provider: {self.llm_provider}. Valid options: {valid_providers}") @@ -155,7 +156,7 @@ def _validate(self) -> None: raise ValueError(f"Invalid max_file_size_mb: {self.max_file_size_mb}. Must be positive.") # Validate severity levels - valid_severities = {'CRITICAL', 'HIGH', 'MEDIUM', 'LOW', 'UNKNOWN'} + valid_severities = Severity.values() severity_list = [s.strip() for s in self.default_severity.split(',')] for severity in severity_list: if severity not in valid_severities: @@ -209,30 +210,30 @@ def get_api_key_for_provider(self) -> str: Raises: EnvironmentError: If API key is not set for the provider """ - if self.llm_provider == "openai": + if self.llm_provider == LLMProvider.OPENAI: if not self.openai_api_key: raise EnvironmentError( "OpenAI API key not found. Set OPENAI_API_KEY environment variable or use --scan-only mode." ) return self.openai_api_key - - elif self.llm_provider == "anthropic": + + elif self.llm_provider == LLMProvider.ANTHROPIC: if not self.anthropic_api_key: raise EnvironmentError( "Anthropic API key not found. Set ANTHROPIC_API_KEY environment variable or use --scan-only mode." ) return self.anthropic_api_key - - elif self.llm_provider == "google": + + elif self.llm_provider == LLMProvider.GOOGLE: if not self.google_api_key: raise EnvironmentError( "Google API key not found. Set GOOGLE_API_KEY environment variable or use --scan-only mode." ) return self.google_api_key - - elif self.llm_provider == "ollama": + + elif self.llm_provider == LLMProvider.OLLAMA: return "" - + else: raise ValueError(f"Unsupported LLM provider: {self.llm_provider}") diff --git a/docksec/docker_scanner.py b/docksec/docker_scanner.py index 0f68c67..3e239b8 100644 --- a/docksec/docker_scanner.py +++ b/docksec/docker_scanner.py @@ -10,6 +10,7 @@ import re from pathlib import Path from docksec.config import RESULTS_DIR, docker_score_prompt +from docksec.enums import Severity from docksec.utils import ScoreResponse, get_llm, print_section, get_custom_logger # Initialize logger @@ -94,9 +95,9 @@ def _validate_severity(severity: str) -> str: if not severity: raise ValueError("Severity cannot be empty") - valid_severities = ['CRITICAL', 'HIGH', 'MEDIUM', 'LOW', 'UNKNOWN'] + valid_severities = Severity.values() severity_list = [s.strip().upper() for s in severity.split(',')] - + for sev in severity_list: if sev not in valid_severities: raise ValueError(f"Invalid severity level: {sev}. Valid values: {', '.join(valid_severities)}") @@ -845,9 +846,9 @@ def add_section_header(self, title): pdf.cell(0, 7, 'No vulnerabilities found.', 0, 1) else: # Count vulnerabilities by severity - severity_counts = {} + severity_counts: Dict[str, int] = {} for vuln in vulnerabilities: - severity = vuln.get('Severity', 'UNKNOWN') + severity = vuln.get('Severity', Severity.UNKNOWN) severity_counts[severity] = severity_counts.get(severity, 0) + 1 pdf.set_font('Arial', '', 10) @@ -986,11 +987,16 @@ def _calculate_local_score(self, results: Dict) -> float: if not vulnerabilities: vuln_score = 100.0 else: - critical = sum(1 for v in vulnerabilities if v.get('Severity') == 'CRITICAL') - high = sum(1 for v in vulnerabilities if v.get('Severity') == 'HIGH') - medium = sum(1 for v in vulnerabilities if v.get('Severity') == 'MEDIUM') - low = sum(1 for v in vulnerabilities if v.get('Severity') == 'LOW') - deduction = (critical * 10) + (high * 5) + (medium * 2) + (low * 1) + severity_weights = { + Severity.CRITICAL: 10, + Severity.HIGH: 5, + Severity.MEDIUM: 2, + Severity.LOW: 1, + } + deduction = sum( + weight * sum(1 for v in vulnerabilities if v.get('Severity') == sev) + for sev, weight in severity_weights.items() + ) vuln_score = max(0.0, 100.0 - deduction) # Configuration score — static Dockerfile checks @@ -1197,9 +1203,9 @@ def _prepare_html_template_vars(self, results: Dict) -> Dict[str, str]: template_vars['DETAILED_VULNERABILITIES_SECTION'] = "" else: # Count vulnerabilities by severity - severity_counts = {'CRITICAL': 0, 'HIGH': 0, 'MEDIUM': 0, 'LOW': 0} + severity_counts = {s: 0 for s in Severity.scored_levels()} for vuln in vulnerabilities: - severity = vuln.get('Severity', 'UNKNOWN') + severity = vuln.get('Severity', Severity.UNKNOWN) if severity in severity_counts: severity_counts[severity] += 1 @@ -1207,19 +1213,19 @@ def _prepare_html_template_vars(self, results: Dict) -> Dict[str, str]: severity_html = f"""
-
{severity_counts['CRITICAL']}
+
{severity_counts[Severity.CRITICAL]}
Critical
-
{severity_counts['HIGH']}
+
{severity_counts[Severity.HIGH]}
High
-
{severity_counts['MEDIUM']}
+
{severity_counts[Severity.MEDIUM]}
Medium
-
{severity_counts['LOW']}
+
{severity_counts[Severity.LOW]}
Low
diff --git a/docksec/enums.py b/docksec/enums.py new file mode 100644 index 0000000..f61d0fd --- /dev/null +++ b/docksec/enums.py @@ -0,0 +1,29 @@ +from enum import Enum + + +class Severity(str, Enum): + CRITICAL = "CRITICAL" + HIGH = "HIGH" + MEDIUM = "MEDIUM" + LOW = "LOW" + UNKNOWN = "UNKNOWN" + + @classmethod + def values(cls) -> list: + return [e.value for e in cls] + + @classmethod + def scored_levels(cls) -> list: + """Severities that affect the security score.""" + return [cls.CRITICAL, cls.HIGH, cls.MEDIUM, cls.LOW] + + +class LLMProvider(str, Enum): + OPENAI = "openai" + ANTHROPIC = "anthropic" + GOOGLE = "google" + OLLAMA = "ollama" + + @classmethod + def values(cls) -> list: + return [e.value for e in cls] diff --git a/docksec/score_calculator.py b/docksec/score_calculator.py index ee37a51..eaf24a3 100644 --- a/docksec/score_calculator.py +++ b/docksec/score_calculator.py @@ -7,6 +7,7 @@ from typing import Dict from docksec.config import docker_score_prompt +from docksec.enums import Severity from docksec.utils import ScoreResponse, get_llm, get_custom_logger # Initialize logger @@ -128,14 +129,16 @@ def get_score_breakdown(self, results: Dict) -> Dict[str, float]: if not vulnerabilities: breakdown['vulnerabilities'] = 100.0 else: - # Weighted by severity - critical_count = sum(1 for v in vulnerabilities if v.get('Severity') == 'CRITICAL') - high_count = sum(1 for v in vulnerabilities if v.get('Severity') == 'HIGH') - medium_count = sum(1 for v in vulnerabilities if v.get('Severity') == 'MEDIUM') - low_count = sum(1 for v in vulnerabilities if v.get('Severity') == 'LOW') - - # Simplified scoring: deduct more for higher severity - deduction = (critical_count * 10) + (high_count * 5) + (medium_count * 2) + (low_count * 1) + severity_weights = { + Severity.CRITICAL: 10, + Severity.HIGH: 5, + Severity.MEDIUM: 2, + Severity.LOW: 1, + } + deduction = sum( + weight * sum(1 for v in vulnerabilities if v.get('Severity') == sev) + for sev, weight in severity_weights.items() + ) breakdown['vulnerabilities'] = max(0, 100 - deduction) # Configuration score derived from actual Dockerfile analysis diff --git a/docksec/utils.py b/docksec/utils.py index 9ea3ed7..dc9dcd7 100644 --- a/docksec/utils.py +++ b/docksec/utils.py @@ -34,6 +34,7 @@ from docksec.config import ( BASE_DIR ) +from docksec.enums import LLMProvider try: from pydantic import BaseModel, Field except ImportError: @@ -162,11 +163,11 @@ def get_llm() -> Union[ChatOpenAI, 'ChatAnthropic', 'ChatGoogleGenerativeAI', 'C logger.info(f"Initializing LLM with provider: {provider}, model: {model}") - if provider == "openai": + if provider == LLMProvider.OPENAI: api_key = config.get_api_key_for_provider() if not os.getenv("OPENAI_API_KEY"): os.environ["OPENAI_API_KEY"] = api_key - + llm = ChatOpenAI( model=model, temperature=temperature, @@ -175,8 +176,8 @@ def get_llm() -> Union[ChatOpenAI, 'ChatAnthropic', 'ChatGoogleGenerativeAI', 'C ) logger.info("OpenAI LLM initialized successfully") return llm - - elif provider == "anthropic": + + elif provider == LLMProvider.ANTHROPIC: if not ANTHROPIC_AVAILABLE: raise ImportError( "Anthropic provider requested but langchain-anthropic is not installed. " @@ -185,7 +186,7 @@ def get_llm() -> Union[ChatOpenAI, 'ChatAnthropic', 'ChatGoogleGenerativeAI', 'C api_key = config.get_api_key_for_provider() if not os.getenv("ANTHROPIC_API_KEY"): os.environ["ANTHROPIC_API_KEY"] = api_key - + llm = ChatAnthropic( model=model, temperature=temperature, @@ -194,8 +195,8 @@ def get_llm() -> Union[ChatOpenAI, 'ChatAnthropic', 'ChatGoogleGenerativeAI', 'C ) logger.info("Anthropic Claude LLM initialized successfully") return llm - - elif provider == "google": + + elif provider == LLMProvider.GOOGLE: if not GOOGLE_AVAILABLE: raise ImportError( "Google provider requested but langchain-google-genai is not installed. " @@ -204,7 +205,7 @@ def get_llm() -> Union[ChatOpenAI, 'ChatAnthropic', 'ChatGoogleGenerativeAI', 'C api_key = config.get_api_key_for_provider() if not os.getenv("GOOGLE_API_KEY"): os.environ["GOOGLE_API_KEY"] = api_key - + llm = ChatGoogleGenerativeAI( model=model, temperature=temperature, @@ -213,8 +214,8 @@ def get_llm() -> Union[ChatOpenAI, 'ChatAnthropic', 'ChatGoogleGenerativeAI', 'C ) logger.info("Google Gemini LLM initialized successfully") return llm - - elif provider == "ollama": + + elif provider == LLMProvider.OLLAMA: if not OLLAMA_AVAILABLE: raise ImportError( "Ollama provider requested but langchain-ollama is not installed. " @@ -228,9 +229,9 @@ def get_llm() -> Union[ChatOpenAI, 'ChatAnthropic', 'ChatGoogleGenerativeAI', 'C ) logger.info(f"Ollama LLM initialized successfully with base URL: {config.ollama_base_url}") return llm - + else: - raise ValueError(f"Unsupported LLM provider: {provider}. Supported: openai, anthropic, google, ollama") + raise ValueError(f"Unsupported LLM provider: {provider}. Supported: {LLMProvider.values()}") except Exception as e: logger.error(f"Failed to initialize LLM: {str(e)}") diff --git a/tests/test_docker_scanner.py b/tests/test_docker_scanner.py index ed062b4..098c9ef 100644 --- a/tests/test_docker_scanner.py +++ b/tests/test_docker_scanner.py @@ -85,10 +85,10 @@ def test_validate_file_path(self): def test_validate_severity(self): """Test severity validation.""" from docksec.docker_scanner import DockerSecurityScanner - + from docksec.enums import Severity + # Valid severities - valid_severities = ["CRITICAL", "HIGH", "MEDIUM", "LOW", "UNKNOWN"] - for sev in valid_severities: + for sev in Severity.values(): result = DockerSecurityScanner._validate_severity(sev) self.assertIn(sev.upper(), result)