diff --git a/docker_scanner.py b/docker_scanner.py index f8d919e..97d9bd8 100644 --- a/docker_scanner.py +++ b/docker_scanner.py @@ -3,6 +3,7 @@ import subprocess import csv import pandas as pd +import logging from typing import List, Tuple, Dict, Optional from datetime import datetime from fpdf import FPDF @@ -12,7 +13,10 @@ from pathlib import Path from config import RESULTS_DIR from config import docker_score_prompt -from utils import ScoreResponse, get_llm, print_section +from utils import ScoreResponse, get_llm, print_section, get_custom_logger + +# Initialize logger +logger = get_custom_logger(__name__) class DockerSecurityScanner: @staticmethod @@ -134,7 +138,11 @@ def __init__(self, dockerfile_path: str, image_name: str, results_dir: str = RES # Verify required tools missing_tools = self._check_tools() if missing_tools: - raise ValueError(f"Missing required tools: {', '.join(missing_tools)}") + error_msg = f"Missing required tools: {', '.join(missing_tools)}\n\n" + error_msg += "Installation instructions:\n" + for tool in missing_tools: + error_msg += f"\n{tool.upper()}:\n{self._get_tool_installation_instructions(tool)}\n" + raise ValueError(error_msg) # Verify Dockerfile exists (after validation) if self.dockerfile_path and not os.path.exists(self.dockerfile_path): @@ -180,6 +188,7 @@ def run_image_only_scan(self, severity: str = "CRITICAL,HIGH") -> Dict: """ # Validate severity input severity = self._validate_severity(severity) + logger.info(f"Starting image-only scan for {self.image_name}") print(f"\n=== Starting image-only scan for {self.image_name} ===") results = { @@ -237,6 +246,31 @@ def _check_tools(self) -> List[str]: missing_tools.append(tool) return missing_tools + + def _get_tool_installation_instructions(self, tool: str) -> str: + """Get installation instructions for a missing tool.""" + instructions = { + 'docker': ( + "Docker is required for image scanning. Please install Docker:\n" + " - Linux: https://docs.docker.com/engine/install/\n" + " - macOS: https://docs.docker.com/desktop/install/mac-install/\n" + " - Windows: https://docs.docker.com/desktop/install/windows-install/" + ), + 'trivy': ( + "Trivy is required for vulnerability scanning. Install it:\n" + " - Linux/Mac: curl -sfL https://raw.githubusercontent.com/aquasecurity/trivy/main/contrib/install.sh | sh -s -- -b /usr/local/bin\n" + " - Windows: See https://aquasecurity.github.io/trivy/latest/getting-started/installation/\n" + " - Or run: python setup_external_tools.py" + ), + 'hadolint': ( + "Hadolint is required for Dockerfile linting. Install it:\n" + " - Linux: curl -L -o hadolint https://github.com/hadolint/hadolint/releases/latest/download/hadolint-Linux-x86_64 && chmod +x hadolint && sudo mv hadolint /usr/local/bin/\n" + " - macOS: brew install hadolint\n" + " - Windows: See https://github.com/hadolint/hadolint#install\n" + " - Or run: python setup_external_tools.py" + ) + } + return instructions.get(tool, f"Please install {tool} from its official documentation.") def scan_dockerfile(self) -> Tuple[bool, Optional[str]]: """ @@ -247,6 +281,7 @@ def scan_dockerfile(self) -> Tuple[bool, Optional[str]]: - bool: True if no issues found, False otherwise - Optional[str]: Output from the scan or None if successful """ + logger.info(f"Starting Dockerfile scan with Hadolint: {self.dockerfile_path}") print("\n=== Starting Dockerfile scan with Hadolint ===") try: result = subprocess.run( @@ -259,16 +294,25 @@ def scan_dockerfile(self) -> Tuple[bool, Optional[str]]: if result.returncode != 0: output = result.stdout if result.stdout else result.stderr + logger.warning(f"Hadolint found issues in {self.dockerfile_path}") print("Dockerfile linting issues found:") print(output) return False, output else: + logger.info("No Dockerfile linting issues found.") print("No Dockerfile linting issues found.") return True, None except subprocess.CalledProcessError as e: + logger.error(f"Error running Hadolint: {e}", exc_info=True) print(f"Error running Hadolint: {e}") return False, str(e) + except subprocess.TimeoutExpired: + logger.error(f"Hadolint scan timed out after 300 seconds for {self.dockerfile_path}") + return False, "Scan timeout" + except Exception as e: + logger.error(f"Unexpected error during Hadolint scan: {e}", exc_info=True) + return False, str(e) def _filter_scan_results(self, scan_results: Dict) -> List[Dict]: """ @@ -321,6 +365,7 @@ def scan_image_json(self, severity: str = "CRITICAL,HIGH") -> Tuple[bool, Option """ # Validate severity input severity = self._validate_severity(severity) + logger.info(f"Starting Trivy JSON scan for image: {self.image_name}") print("\n=== Starting vulnerability scan with Trivy for Json Output ===") try: @@ -376,6 +421,7 @@ def scan_image(self, severity: str = "CRITICAL,HIGH") -> Tuple[bool, Optional[st """ # Validate severity input severity = self._validate_severity(severity) + logger.info(f"Starting Trivy scan for image: {self.image_name} with severity: {severity}") print("\n=== Starting vulnerability scan with Trivy ===") try: diff --git a/pyproject.toml b/pyproject.toml index 6cc6413..2333768 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta" [tool.black] line-length = 88 -target-version = ['py38'] +target-version = ['py312'] include = '\.pyi?$' [tool.isort] diff --git a/requirements.txt b/requirements.txt index 8885173..8e914a6 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,13 +1,18 @@ -pydantic -langchain-core -langchain -langchain-openai -python-dotenv -pandas +# Core dependencies +pydantic>=2.0.0,<3.0.0 +langchain-core>=0.3.0,<2.0.0 +langchain>=0.3.0,<2.0.0 +langchain-openai>=0.2.0,<1.0.0 +python-dotenv>=1.0.0,<2.0.0 +pandas>=2.0.0,<3.0.0 -tqdm -colorama -rich +# UI and progress +tqdm>=4.65.0,<5.0.0 +colorama>=0.4.6,<1.0.0 +rich>=13.0.0,<14.0.0 -fpdf -setuptools \ No newline at end of file +# PDF generation +fpdf2>=2.7.0,<3.0.0 + +# Build tools +setuptools>=65.0.0 \ No newline at end of file diff --git a/setup.py b/setup.py index 6ef50a2..2d16aba 100644 --- a/setup.py +++ b/setup.py @@ -22,18 +22,17 @@ }, python_requires=">=3.12", install_requires=[ - "pydantic", - "langchain-core", - "langchain", - "langchain-openai", - "python-dotenv", - "pandas", - - "tqdm", - "colorama", - "rich", - "fpdf", - "setuptools", + "pydantic>=2.0.0,<3.0.0", + "langchain-core>=0.3.0,<2.0.0", + "langchain>=0.3.0,<2.0.0", + "langchain-openai>=0.2.0,<1.0.0", + "python-dotenv>=1.0.0,<2.0.0", + "pandas>=2.0.0,<3.0.0", + "tqdm>=4.65.0,<5.0.0", + "colorama>=0.4.6,<1.0.0", + "rich>=13.0.0,<14.0.0", + "fpdf2>=2.7.0,<3.0.0", + "setuptools>=65.0.0", ], classifiers=[ "Programming Language :: Python :: 3", diff --git a/tests/README.md b/tests/README.md new file mode 100644 index 0000000..1c1e237 --- /dev/null +++ b/tests/README.md @@ -0,0 +1,46 @@ +# DockSec Test Suite + +This directory contains unit and integration tests for DockSec. + +## Test Structure + +- `test_docker_scanner.py`: Unit tests for the DockerSecurityScanner class +- `test_utils.py`: Unit tests for utility functions +- `test_config.py`: Unit tests for configuration management +- `test_integration.py`: Integration tests that test multiple components together +- `conftest.py`: Pytest fixtures and configuration + +## Running Tests + +### Using unittest +```bash +python -m unittest discover tests +``` + +### Using pytest (if installed) +```bash +pytest tests/ +``` + +### Running specific test file +```bash +python -m unittest tests.test_docker_scanner +``` + +## Test Coverage + +Currently, tests cover: +- Input validation (image names, file paths, severity levels) +- Tool checking and installation instructions +- Dockerfile loading +- Configuration and API key handling +- LLM initialization +- Basic scanner initialization + +## Adding New Tests + +When adding new functionality, please add corresponding tests: +1. Unit tests for individual functions/methods +2. Integration tests for workflows +3. Update this README if adding new test categories + diff --git a/tests/__init__.py b/tests/__init__.py new file mode 100644 index 0000000..0c01286 --- /dev/null +++ b/tests/__init__.py @@ -0,0 +1,2 @@ +"""Tests for DockSec package.""" + diff --git a/tests/conftest.py b/tests/conftest.py new file mode 100644 index 0000000..7fe3741 --- /dev/null +++ b/tests/conftest.py @@ -0,0 +1,23 @@ +"""Pytest configuration and fixtures.""" +import pytest +import os +import tempfile +import shutil + + +@pytest.fixture +def temp_dir(): + """Create a temporary directory for tests.""" + test_dir = tempfile.mkdtemp() + yield test_dir + shutil.rmtree(test_dir) + + +@pytest.fixture +def sample_dockerfile(temp_dir): + """Create a sample Dockerfile for testing.""" + dockerfile_path = os.path.join(temp_dir, "Dockerfile") + with open(dockerfile_path, 'w') as f: + f.write("FROM ubuntu:latest\nRUN echo 'test'\n") + return dockerfile_path + diff --git a/tests/test_config.py b/tests/test_config.py new file mode 100644 index 0000000..c377544 --- /dev/null +++ b/tests/test_config.py @@ -0,0 +1,38 @@ +"""Unit tests for configuration.""" +import unittest +import os +from unittest.mock import patch + + +class TestConfig(unittest.TestCase): + """Test cases for configuration.""" + + @patch.dict(os.environ, {}, clear=True) + def test_get_openai_api_key_missing(self): + """Test API key retrieval when not set.""" + from config import get_openai_api_key + + with self.assertRaises(EnvironmentError): + get_openai_api_key() + + @patch.dict(os.environ, {'OPENAI_API_KEY': 'test-key-123'}) + def test_get_openai_api_key_present(self): + """Test API key retrieval when set.""" + from config import get_openai_api_key + + api_key = get_openai_api_key() + self.assertEqual(api_key, 'test-key-123') + + def test_prompt_templates_exist(self): + """Test that prompt templates are defined.""" + from config import docker_agent_template, docker_score_template + + self.assertIsNotNone(docker_agent_template) + self.assertIsNotNone(docker_score_template) + self.assertIn("Dockerfile", docker_agent_template) + self.assertIn("score", docker_score_template.lower()) + + +if __name__ == '__main__': + unittest.main() + diff --git a/tests/test_docker_scanner.py b/tests/test_docker_scanner.py new file mode 100644 index 0000000..934a19f --- /dev/null +++ b/tests/test_docker_scanner.py @@ -0,0 +1,154 @@ +"""Unit tests for DockerSecurityScanner class.""" +import unittest +import os +import tempfile +from unittest.mock import Mock, patch, MagicMock +from pathlib import Path + +# Import after mocking external dependencies +import sys +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + + +class TestDockerSecurityScanner(unittest.TestCase): + """Test cases for DockerSecurityScanner.""" + + def setUp(self): + """Set up test fixtures.""" + self.test_dockerfile = None + self.test_dir = None + + def tearDown(self): + """Clean up test fixtures.""" + if self.test_dir and os.path.exists(self.test_dir): + import shutil + shutil.rmtree(self.test_dir) + + def create_test_dockerfile(self, content="FROM ubuntu:latest"): + """Create a temporary Dockerfile for testing.""" + self.test_dir = tempfile.mkdtemp() + self.test_dockerfile = os.path.join(self.test_dir, "Dockerfile") + with open(self.test_dockerfile, 'w') as f: + f.write(content) + return self.test_dockerfile + + @patch('docker_scanner.subprocess.run') + @patch('docker_scanner.get_llm') + def test_init_with_valid_inputs(self, mock_llm, mock_subprocess): + """Test initialization with valid inputs.""" + # Mock subprocess calls for tool checking and docker image inspect + mock_subprocess.return_value = Mock(returncode=0, stdout="", stderr="") + + # Mock LLM + mock_llm.return_value = Mock() + + dockerfile = self.create_test_dockerfile() + + from docker_scanner import DockerSecurityScanner + + scanner = DockerSecurityScanner(dockerfile, "test:latest") + self.assertEqual(scanner.dockerfile_path, dockerfile) + self.assertEqual(scanner.image_name, "test:latest") + self.assertIsNone(scanner.analysis_score) + + def test_validate_image_name(self): + """Test image name validation.""" + from docker_scanner import DockerSecurityScanner + + # Valid image names + valid_names = ["nginx:latest", "myimage:v1.0", "registry/image:tag"] + for name in valid_names: + result = DockerSecurityScanner._validate_image_name(name) + self.assertEqual(result, name) + + # Invalid image names + invalid_names = ["", "../../etc/passwd", "image with spaces", "image\nnewline"] + for name in invalid_names: + with self.assertRaises(ValueError): + DockerSecurityScanner._validate_image_name(name) + + def test_validate_file_path(self): + """Test file path validation.""" + from docker_scanner import DockerSecurityScanner + + # Path traversal attempts should be rejected + with self.assertRaises(ValueError): + DockerSecurityScanner._validate_file_path("../../../etc/passwd") + + # Valid path should work + dockerfile = self.create_test_dockerfile() + result = DockerSecurityScanner._validate_file_path(dockerfile) + self.assertTrue(result.exists()) + + def test_validate_severity(self): + """Test severity validation.""" + from docker_scanner import DockerSecurityScanner + + # Valid severities + valid_severities = ["CRITICAL", "HIGH", "MEDIUM", "LOW", "UNKNOWN"] + for sev in valid_severities: + result = DockerSecurityScanner._validate_severity(sev) + self.assertIn(sev.upper(), result) + + # Invalid severity + with self.assertRaises(ValueError): + DockerSecurityScanner._validate_severity("INVALID") + + # Multiple valid severities + result = DockerSecurityScanner._validate_severity("CRITICAL,HIGH") + self.assertIn("CRITICAL", result) + self.assertIn("HIGH", result) + + @patch('docker_scanner.subprocess.run') + def test_check_tools_missing(self, mock_subprocess): + """Test tool checking with missing tools.""" + from docker_scanner import DockerSecurityScanner + + # Mock FileNotFoundError for missing tool + mock_subprocess.side_effect = FileNotFoundError() + + dockerfile = self.create_test_dockerfile() + + with patch('docker_scanner.get_llm'): + scanner = DockerSecurityScanner.__new__(DockerSecurityScanner) + scanner.required_tools = ['docker', 'trivy'] + missing = scanner._check_tools() + self.assertEqual(missing, ['docker', 'trivy']) + + @patch('docker_scanner.subprocess.run') + def test_check_tools_present(self, mock_subprocess): + """Test tool checking with all tools present.""" + from docker_scanner import DockerSecurityScanner + + # Mock successful tool check + mock_subprocess.return_value = Mock(returncode=0, stdout="", stderr="") + + scanner = DockerSecurityScanner.__new__(DockerSecurityScanner) + scanner.required_tools = ['docker', 'trivy'] + missing = scanner._check_tools() + self.assertEqual(missing, []) + + def test_get_tool_installation_instructions(self): + """Test installation instructions for tools.""" + from docker_scanner import DockerSecurityScanner + + scanner = DockerSecurityScanner.__new__(DockerSecurityScanner) + + # Test known tools + docker_instructions = scanner._get_tool_installation_instructions('docker') + self.assertIn('Docker', docker_instructions) + + trivy_instructions = scanner._get_tool_installation_instructions('trivy') + self.assertIn('Trivy', trivy_instructions) + + hadolint_instructions = scanner._get_tool_installation_instructions('hadolint') + self.assertIn('Hadolint', hadolint_instructions) + + # Test unknown tool + unknown_instructions = scanner._get_tool_installation_instructions('unknown') + self.assertIn('unknown', unknown_instructions) + + +if __name__ == '__main__': + unittest.main() + diff --git a/tests/test_integration.py b/tests/test_integration.py new file mode 100644 index 0000000..d18dce9 --- /dev/null +++ b/tests/test_integration.py @@ -0,0 +1,93 @@ +"""Integration tests for DockSec.""" +import unittest +import os +import tempfile +import shutil +from unittest.mock import patch, Mock + +# Import after mocking external dependencies +import sys +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + + +class TestIntegration(unittest.TestCase): + """Integration tests that test multiple components together.""" + + def setUp(self): + """Set up test fixtures.""" + self.test_dir = tempfile.mkdtemp() + self.test_dockerfile = os.path.join(self.test_dir, "Dockerfile") + with open(self.test_dockerfile, 'w') as f: + f.write("FROM ubuntu:latest\nRUN echo 'test'") + + def tearDown(self): + """Clean up test fixtures.""" + if os.path.exists(self.test_dir): + shutil.rmtree(self.test_dir) + + @patch('docker_scanner.subprocess.run') + @patch('docker_scanner.get_llm') + @patch('config.get_openai_api_key') + def test_full_scan_workflow(self, mock_api_key, mock_llm, mock_subprocess): + """Test complete scanning workflow.""" + # Mock API key + mock_api_key.return_value = "test-api-key" + + # Mock LLM responses + mock_llm_instance = Mock() + mock_llm.return_value = mock_llm_instance + + # Mock all subprocess calls + def subprocess_side_effect(*args, **kwargs): + if 'docker' in args[0] and 'inspect' in args[0]: + return Mock(returncode=0, stdout='{"Id": "test"}', stderr='') + elif '--version' in args[0]: + return Mock(returncode=0, stdout='version 1.0', stderr='') + else: + return Mock(returncode=0, stdout='', stderr='') + + mock_subprocess.side_effect = subprocess_side_effect + + from docker_scanner import DockerSecurityScanner + + scanner = DockerSecurityScanner(self.test_dockerfile, "test:latest") + + # Verify initialization + self.assertIsNotNone(scanner) + self.assertEqual(scanner.dockerfile_path, self.test_dockerfile) + self.assertEqual(scanner.image_name, "test:latest") + + @patch('docker_scanner.subprocess.run') + @patch('docker_scanner.get_llm') + @patch('config.get_openai_api_key') + def test_image_only_scan(self, mock_api_key, mock_llm, mock_subprocess): + """Test image-only scanning without Dockerfile.""" + # Mock API key + mock_api_key.return_value = "test-api-key" + + # Mock LLM + mock_llm_instance = Mock() + mock_llm.return_value = mock_llm_instance + + # Mock subprocess calls + def subprocess_side_effect(*args, **kwargs): + if 'docker' in args[0] and 'inspect' in args[0]: + return Mock(returncode=0, stdout='{"Id": "test"}', stderr='') + elif '--version' in args[0]: + return Mock(returncode=0, stdout='version 1.0', stderr='') + else: + return Mock(returncode=0, stdout='[]', stderr='') + + mock_subprocess.side_effect = subprocess_side_effect + + from docker_scanner import DockerSecurityScanner + + # Should work without Dockerfile + scanner = DockerSecurityScanner(None, "test:latest") + self.assertIsNone(scanner.dockerfile_path) + self.assertEqual(scanner.image_name, "test:latest") + + +if __name__ == '__main__': + unittest.main() + diff --git a/tests/test_utils.py b/tests/test_utils.py new file mode 100644 index 0000000..2eb8426 --- /dev/null +++ b/tests/test_utils.py @@ -0,0 +1,74 @@ +"""Unit tests for utility functions.""" +import unittest +import os +import tempfile +from unittest.mock import patch, Mock + +# Import after mocking external dependencies +import sys +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + + +class TestUtils(unittest.TestCase): + """Test cases for utility functions.""" + + def test_get_custom_logger(self): + """Test logger creation.""" + from utils import get_custom_logger + + logger = get_custom_logger('TestLogger') + self.assertEqual(logger.name, 'TestLogger') + self.assertEqual(logger.level, 20) # INFO level + + def test_load_docker_file(self): + """Test Dockerfile loading.""" + from utils import load_docker_file + + # Create temporary Dockerfile + with tempfile.NamedTemporaryFile(mode='w', suffix='.dockerfile', delete=False) as f: + f.write("FROM ubuntu:latest\nRUN echo 'test'") + temp_path = f.name + + try: + content = load_docker_file(temp_path) + self.assertIn("FROM ubuntu:latest", content) + self.assertIn("RUN echo 'test'", content) + finally: + os.unlink(temp_path) + + def test_load_docker_file_not_found(self): + """Test Dockerfile loading when file doesn't exist.""" + from utils import load_docker_file + + result = load_docker_file("/nonexistent/path/Dockerfile") + self.assertIsNone(result) + + @patch('utils.get_openai_api_key') + @patch('utils.ChatOpenAI') + def test_get_llm(self, mock_chatopenai, mock_api_key): + """Test LLM initialization.""" + from utils import get_llm + + mock_api_key.return_value = "test-api-key" + mock_llm_instance = Mock() + mock_chatopenai.return_value = mock_llm_instance + + llm = get_llm() + + mock_chatopenai.assert_called_once() + self.assertIsNotNone(llm) + + @patch('utils.get_openai_api_key') + def test_get_llm_no_api_key(self, mock_api_key): + """Test LLM initialization without API key.""" + from utils import get_llm + + mock_api_key.side_effect = EnvironmentError("API key not found") + + with self.assertRaises(EnvironmentError): + get_llm() + + +if __name__ == '__main__': + unittest.main() +