refactor: replace magic strings with Severity and LLMProvider enums (#47)

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()}')

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}")
 

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,29 +1203,29 @@ 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
 
             # Create severity statistics HTML
             severity_html = f"""
             <div class="severity-stats">
                 <div class="severity-item severity-critical">
-                    <div class="severity-count">{severity_counts['CRITICAL']}</div>
+                    <div class="severity-count">{severity_counts[Severity.CRITICAL]}</div>
                     <div class="severity-label">Critical</div>
                 </div>
                 <div class="severity-item severity-high">
-                    <div class="severity-count">{severity_counts['HIGH']}</div>
+                    <div class="severity-count">{severity_counts[Severity.HIGH]}</div>
                     <div class="severity-label">High</div>
                 </div>
                 <div class="severity-item severity-medium">
-                    <div class="severity-count">{severity_counts['MEDIUM']}</div>
+                    <div class="severity-count">{severity_counts[Severity.MEDIUM]}</div>
                     <div class="severity-label">Medium</div>
                 </div>
                 <div class="severity-item severity-low">
-                    <div class="severity-count">{severity_counts['LOW']}</div>
+                    <div class="severity-count">{severity_counts[Severity.LOW]}</div>
                     <div class="severity-label">Low</div>
                 </div>
             </div>

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]

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

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)}")

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)