Initial commit: AASRT v1.0.0 - AI Agent Security Reconnaissance Tool

dev/tests/integration/test_scan_workflow.py

@@ -0,0 +1,198 @@
+"""
+Integration Tests for Scan Workflow
+
+Tests the complete scan workflow from query to report generation.
+"""
+
+import pytest
+import tempfile
+from pathlib import Path
+from unittest.mock import MagicMock, patch
+from datetime import datetime
+
+
+class TestEndToEndScan:
+    """Integration tests for complete scan workflow."""
+
+    @pytest.fixture
+    def mock_shodan_response(self, sample_search_results):
+        """Mock Shodan API response."""
+        return {
+            'matches': sample_search_results,
+            'total': len(sample_search_results)
+        }
+
+    @pytest.fixture
+    def temp_workspace(self):
+        """Create temporary workspace for test files."""
+        with tempfile.TemporaryDirectory() as tmpdir:
+            workspace = Path(tmpdir)
+            (workspace / 'reports').mkdir()
+            (workspace / 'data').mkdir()
+            (workspace / 'logs').mkdir()
+            yield workspace
+
+    def test_scan_template_workflow(self, mock_shodan_response, temp_workspace):
+        """Test scanning using a template."""
+        from unittest.mock import patch, MagicMock
+        
+        mock_client = MagicMock()
+        mock_client.info.return_value = {'plan': 'dev', 'query_credits': 100}
+        mock_client.search.return_value = mock_shodan_response
+        
+        with patch('shodan.Shodan', return_value=mock_client):
+            with patch.dict('os.environ', {
+                'SHODAN_API_KEY': 'test_key_12345',
+                'AASRT_REPORTS_DIR': str(temp_workspace / 'reports'),
+                'AASRT_DATA_DIR': str(temp_workspace / 'data'),
+            }):
+                # Import after patching
+                from src.core.query_manager import QueryManager
+                from src.utils.config import Config
+                
+                config = Config()
+                qm = QueryManager(config)
+                
+                # Check templates are available
+                templates = qm.get_available_templates()
+                assert len(templates) > 0
+
+    def test_custom_query_workflow(self, mock_shodan_response, temp_workspace):
+        """Test scanning with a custom query."""
+        mock_client = MagicMock()
+        mock_client.info.return_value = {'plan': 'dev', 'query_credits': 100}
+        mock_client.search.return_value = mock_shodan_response
+        
+        with patch('shodan.Shodan', return_value=mock_client):
+            with patch.dict('os.environ', {
+                'SHODAN_API_KEY': 'test_key_12345',
+            }):
+                from src.engines.shodan_engine import ShodanEngine
+                from src.utils.config import Config
+                
+                config = Config()
+                engine = ShodanEngine(config=config)
+                engine._client = mock_client
+                
+                results = engine.search('http.title:"Test"')
+                assert len(results) > 0
+
+
+class TestVulnerabilityAssessmentIntegration:
+    """Integration tests for vulnerability assessment pipeline."""
+
+    def test_assess_search_results(self, sample_search_results):
+        """Test vulnerability assessment on search results."""
+        from src.core.vulnerability_assessor import VulnerabilityAssessor
+        from src.engines.base import SearchResult
+        
+        assessor = VulnerabilityAssessor()
+        
+        # Convert sample data to SearchResult
+        result = SearchResult(
+            ip=sample_search_results[0]['ip_str'],
+            port=sample_search_results[0]['port'],
+            protocol='tcp',
+            banner=sample_search_results[0].get('data', ''),
+            metadata=sample_search_results[0]
+        )
+        
+        vulns = assessor.assess(result)
+        # Should return a list (may be empty if no vulns detected)
+        assert isinstance(vulns, list)
+
+    def test_risk_scoring_integration(self, sample_search_results):
+        """Test risk scoring on assessed results."""
+        from src.core.risk_scorer import RiskScorer
+        from src.core.vulnerability_assessor import VulnerabilityAssessor
+        from src.engines.base import SearchResult
+        
+        assessor = VulnerabilityAssessor()
+        scorer = RiskScorer()
+        
+        result = SearchResult(
+            ip=sample_search_results[0]['ip_str'],
+            port=sample_search_results[0]['port'],
+            protocol='tcp',
+            banner=sample_search_results[0].get('data', ''),
+            metadata=sample_search_results[0]
+        )
+        
+        vulns = assessor.assess(result)
+        score = scorer.score(result)
+        
+        assert 0 <= score <= 100
+
+
+class TestReportGenerationIntegration:
+    """Integration tests for report generation."""
+
+    @pytest.fixture
+    def temp_reports_dir(self):
+        """Create temporary reports directory."""
+        with tempfile.TemporaryDirectory() as tmpdir:
+            yield Path(tmpdir)
+
+    def test_json_report_generation(self, temp_reports_dir, sample_search_results):
+        """Test JSON report generation."""
+        from src.reporting import JSONReporter, ScanReport
+        from src.engines.base import SearchResult
+        
+        # Create scan report data
+        results = [
+            SearchResult(
+                ip=r['ip_str'],
+                port=r['port'],
+                protocol='tcp',
+                banner=r.get('data', ''),
+                metadata=r
+            ) for r in sample_search_results
+        ]
+        
+        report = ScanReport(
+            scan_id='test-scan-001',
+            query='test query',
+            engine='shodan',
+            started_at=datetime.utcnow(),
+            completed_at=datetime.utcnow(),
+            results=results,
+            total_results=len(results)
+        )
+        
+        reporter = JSONReporter(output_dir=str(temp_reports_dir))
+        output_path = reporter.generate(report)
+        
+        assert Path(output_path).exists()
+        assert output_path.endswith('.json')
+
+    def test_csv_report_generation(self, temp_reports_dir, sample_search_results):
+        """Test CSV report generation."""
+        from src.reporting import CSVReporter, ScanReport
+        from src.engines.base import SearchResult
+        
+        results = [
+            SearchResult(
+                ip=r['ip_str'],
+                port=r['port'],
+                protocol='tcp',
+                banner=r.get('data', ''),
+                metadata=r
+            ) for r in sample_search_results
+        ]
+        
+        report = ScanReport(
+            scan_id='test-scan-002',
+            query='test query',
+            engine='shodan',
+            started_at=datetime.utcnow(),
+            completed_at=datetime.utcnow(),
+            results=results,
+            total_results=len(results)
+        )
+        
+        reporter = CSVReporter(output_dir=str(temp_reports_dir))
+        output_path = reporter.generate(report)
+        
+        assert Path(output_path).exists()
+        assert output_path.endswith('.csv')
+