diff --git a/api/src/backend/api/tests/test_utils.py b/api/src/backend/api/tests/test_utils.py index 352c8262fa..e94504e902 100644 --- a/api/src/backend/api/tests/test_utils.py +++ b/api/src/backend/api/tests/test_utils.py @@ -6,16 +6,18 @@ from rest_framework.exceptions import NotFound, ValidationError from api.db_router import MainRouter from api.exceptions import InvitationTokenExpiredException -from api.models import Invitation, Provider +from api.models import Integration, Invitation, Provider from api.utils import ( get_prowler_provider_kwargs, initialize_prowler_provider, merge_dicts, + prowler_integration_connection_test, prowler_provider_connection_test, return_prowler_provider, validate_invitation, ) from prowler.providers.aws.aws_provider import AwsProvider +from prowler.providers.aws.lib.security_hub.security_hub import SecurityHubConnection from prowler.providers.azure.azure_provider import AzureProvider from prowler.providers.gcp.gcp_provider import GcpProvider from prowler.providers.kubernetes.kubernetes_provider import KubernetesProvider @@ -390,3 +392,109 @@ class TestValidateInvitation: mock_db.get.assert_called_once_with( token="VALID_TOKEN", email__iexact="user@example.com" ) + + +class TestProwlerIntegrationConnectionTest: + """Test prowler_integration_connection_test function for SecurityHub regions reset.""" + + @patch("api.utils.SecurityHub") + def test_security_hub_connection_failure_resets_regions( + self, mock_security_hub_class + ): + """Test that SecurityHub connection failure resets regions to empty dict.""" + # Create integration with existing regions configuration + integration = MagicMock() + integration.integration_type = Integration.IntegrationChoices.AWS_SECURITY_HUB + integration.credentials = { + "aws_access_key_id": "test_key", + "aws_secret_access_key": "test_secret", + } + integration.configuration = { + "send_only_fails": True, + "regions": { + "us-east-1": True, + "us-west-2": True, + "eu-west-1": False, + "ap-south-1": False, + }, + } + + # Mock provider relationship + mock_provider = MagicMock() + mock_provider.uid = "123456789012" + mock_relationship = MagicMock() + mock_relationship.provider = mock_provider + integration.integrationproviderrelationship_set.first.return_value = ( + mock_relationship + ) + + # Mock failed SecurityHub connection + mock_connection = SecurityHubConnection( + is_connected=False, + error=Exception("SecurityHub testing"), + enabled_regions=set(), + disabled_regions=set(), + ) + mock_security_hub_class.test_connection.return_value = mock_connection + + # Call the function + result = prowler_integration_connection_test(integration) + + # Assertions + assert result.is_connected is False + assert str(result.error) == "SecurityHub testing" + + # Verify regions were completely reset to empty dict + assert integration.configuration["regions"] == {} + + # Verify save was called to persist the change + integration.save.assert_called_once() + + # Verify test_connection was called with correct parameters + mock_security_hub_class.test_connection.assert_called_once_with( + aws_account_id="123456789012", + raise_on_exception=False, + aws_access_key_id="test_key", + aws_secret_access_key="test_secret", + ) + + @patch("api.utils.SecurityHub") + def test_security_hub_connection_success_saves_regions( + self, mock_security_hub_class + ): + """Test that successful SecurityHub connection saves regions correctly.""" + integration = MagicMock() + integration.integration_type = Integration.IntegrationChoices.AWS_SECURITY_HUB + integration.credentials = { + "aws_access_key_id": "valid_key", + "aws_secret_access_key": "valid_secret", + } + integration.configuration = {"send_only_fails": False} + + # Mock provider relationship + mock_provider = MagicMock() + mock_provider.uid = "123456789012" + mock_relationship = MagicMock() + mock_relationship.provider = mock_provider + integration.integrationproviderrelationship_set.first.return_value = ( + mock_relationship + ) + + # Mock successful SecurityHub connection with regions + mock_connection = SecurityHubConnection( + is_connected=True, + error=None, + enabled_regions={"us-east-1", "eu-west-1"}, + disabled_regions={"ap-south-1"}, + ) + mock_security_hub_class.test_connection.return_value = mock_connection + + result = prowler_integration_connection_test(integration) + + assert result.is_connected is True + + # Verify regions were saved correctly + assert integration.configuration["regions"]["us-east-1"] is True + assert integration.configuration["regions"]["eu-west-1"] is True + assert integration.configuration["regions"]["ap-south-1"] is False + integration.save.assert_called_once() diff --git a/api/src/backend/api/utils.py b/api/src/backend/api/utils.py index 7cbb7a0247..1cc10674d7 100644 --- a/api/src/backend/api/utils.py +++ b/api/src/backend/api/utils.py @@ -223,6 +223,10 @@ def prowler_integration_connection_test(integration: Integration) -> Connection: # Save regions information in the integration configuration integration.configuration["regions"] = regions_status integration.save() + else: + # Reset regions information if connection fails + integration.configuration["regions"] = {} + integration.save() return connection elif integration.integration_type == Integration.IntegrationChoices.JIRA: diff --git a/api/src/backend/tasks/jobs/integrations.py b/api/src/backend/tasks/jobs/integrations.py index 5cc25d09a8..e95d39652e 100644 --- a/api/src/backend/tasks/jobs/integrations.py +++ b/api/src/backend/tasks/jobs/integrations.py @@ -220,6 +220,11 @@ def get_security_hub_client_from_integration( **credentials, ) return True, security_hub + else: + # Reset regions information if connection fails + with rls_transaction(tenant_id): + integration.configuration["regions"] = {} + integration.save() return False, connection diff --git a/api/src/backend/tasks/tests/test_integrations.py b/api/src/backend/tasks/tests/test_integrations.py index 4dd68b9f95..1a55f14f2c 100644 --- a/api/src/backend/tasks/tests/test_integrations.py +++ b/api/src/backend/tasks/tests/test_integrations.py @@ -492,14 +492,17 @@ class TestProwlerIntegrationConnectionTest: def test_aws_security_hub_integration_connection_failure( self, mock_security_hub_class ): - """Test AWS Security Hub integration connection failure.""" + """Test AWS Security Hub integration connection failure resets regions.""" integration = MagicMock() integration.integration_type = Integration.IntegrationChoices.AWS_SECURITY_HUB integration.credentials = { "aws_access_key_id": "invalid_key", "aws_secret_access_key": "invalid_secret", } - integration.configuration = {"send_only_fails": False} + integration.configuration = { + "send_only_fails": False, + "regions": {"us-east-1": True, "us-west-2": False}, # Existing regions + } # Mock integration provider relationship mock_provider = MagicMock() @@ -524,8 +527,9 @@ class TestProwlerIntegrationConnectionTest: assert result.is_connected is False assert result.error == test_exception - # Verify regions were not saved when connection failed - integration.save.assert_not_called() + # Verify regions were reset to empty dict when connection failed + assert integration.configuration["regions"] == {} + integration.save.assert_called_once() @patch("api.utils.SecurityHub") def test_aws_security_hub_integration_with_provider_credentials( @@ -575,6 +579,72 @@ class TestProwlerIntegrationConnectionTest: assert not integration.configuration["regions"]["ap-south-1"] integration.save.assert_called_once() + @patch("api.utils.SecurityHub") + def test_aws_security_hub_connection_failure_with_multiple_regions_clears_all( + self, mock_security_hub_class + ): + """Test that SecurityHub connection failure clears all existing regions data.""" + integration = MagicMock() + integration.integration_type = Integration.IntegrationChoices.AWS_SECURITY_HUB + integration.credentials = { + "aws_access_key_id": "test_key", + "aws_secret_access_key": "test_secret", + } + # Start with complex regions configuration + integration.configuration = { + "send_only_fails": True, + "regions": { + "us-east-1": True, + "us-east-2": False, + "us-west-1": True, + "us-west-2": True, + "eu-west-1": False, + "eu-west-2": True, + "eu-central-1": True, + "ap-northeast-1": False, + "ap-southeast-1": True, + "ap-southeast-2": False, + }, + } + + # Mock integration provider relationship + mock_provider = MagicMock() + mock_provider.uid = "987654321098" + mock_relationship = MagicMock() + mock_relationship.provider = mock_provider + integration.integrationproviderrelationship_set.first.return_value = ( + mock_relationship + ) + + # Mock failed SecurityHub connection + mock_connection = SecurityHubConnection( + is_connected=False, + error=Exception("Invalid credentials or permissions"), + enabled_regions=set(), + disabled_regions=set(), + ) + mock_security_hub_class.test_connection.return_value = mock_connection + + result = prowler_integration_connection_test(integration) + + assert result.is_connected is False + assert str(result.error) == "Invalid credentials or permissions" + + # Verify all regions were completely cleared + assert integration.configuration["regions"] == {} + assert len(integration.configuration["regions"]) == 0 + + # Verify save was called to persist the cleared regions + integration.save.assert_called_once() + + # Verify the test_connection was called with correct parameters + mock_security_hub_class.test_connection.assert_called_once_with( + aws_account_id="987654321098", + raise_on_exception=False, + aws_access_key_id="test_key", + aws_secret_access_key="test_secret", + ) + def test_unsupported_integration_type(self): """Test unsupported integration type raises ValueError.""" integration = MagicMock() @@ -674,15 +744,19 @@ class TestSecurityHubIntegrationUploads: "us-west-2", ] + @patch("tasks.jobs.integrations.rls_transaction") @patch("tasks.jobs.integrations.SecurityHub.test_connection") @patch("tasks.jobs.integrations.initialize_prowler_provider") def test_get_security_hub_client_from_integration_failure( - self, mock_initialize_provider, mock_test_connection + self, mock_initialize_provider, mock_test_connection, mock_rls ): - """Test SecurityHub client creation failure.""" + """Test SecurityHub client creation failure resets regions.""" # Mock integration mock_integration = MagicMock() - mock_integration.configuration = {"send_only_fails": False} + mock_integration.configuration = { + "send_only_fails": False, + "regions": {"us-east-1": True, "us-west-2": False}, # Existing regions + } mock_integration.credentials = {} # Empty credentials, use provider # Mock tenant_id @@ -710,6 +784,9 @@ class TestSecurityHubIntegrationUploads: # Mock findings mock_findings = [{"finding": "test"}] + # Mock RLS context manager + mock_rls.return_value.__enter__.return_value = None + connected, connection = get_security_hub_client_from_integration( mock_integration, tenant_id, mock_findings ) @@ -725,6 +802,93 @@ class TestSecurityHubIntegrationUploads: aws_secret_access_key="test_secret_key", ) + # Verify regions were reset to empty when connection failed + assert mock_integration.configuration["regions"] == {} + mock_integration.save.assert_called_once() + # Verify RLS transaction was used for the reset + assert ( + mock_rls.call_count == 2 + ) # Once for getting provider, once for resetting regions + + @patch("tasks.jobs.integrations.rls_transaction") + @patch("tasks.jobs.integrations.SecurityHub.test_connection") + def test_get_security_hub_client_from_integration_failure_clears_existing_regions( + self, mock_test_connection, mock_rls + ): + """Test that SecurityHub client creation failure clears existing regions configuration.""" + # Mock integration with pre-existing regions configuration + mock_integration = MagicMock() + mock_integration.configuration = { + "send_only_fails": True, + "regions": { + "us-east-1": True, + "us-west-2": True, + "eu-west-1": False, + "ap-south-1": False, + }, # Pre-existing regions configuration + } + mock_integration.credentials = { + "aws_access_key_id": "test_key_id", + "aws_secret_access_key": "test_secret_key", + } + + # Mock tenant_id + tenant_id = "550e8400-e29b-41d4-a716-446655440000" + + # Mock provider relationship + mock_provider = MagicMock() + mock_provider.uid = "123456789012" + mock_provider.secret.secret = { + "aws_access_key_id": "provider_key", + "aws_secret_access_key": "provider_secret", + } + mock_relationship = MagicMock() + mock_relationship.provider = mock_provider + mock_integration.integrationproviderrelationship_set.first.return_value = ( + mock_relationship + ) + + # Mock failed connection with specific error + mock_connection = MagicMock() + mock_connection.is_connected = False + mock_connection.error = "Access denied: SecurityHub not enabled in region" + mock_test_connection.return_value = mock_connection + + # Mock findings + mock_findings = [{"finding": "test1"}, {"finding": "test2"}] + + # Mock RLS context manager + mock_rls.return_value.__enter__.return_value = None + + # Call the function + connected, connection = get_security_hub_client_from_integration( + mock_integration, tenant_id, mock_findings + ) + + # Assertions + assert connected is False + assert connection == mock_connection + assert connection.error == "Access denied: SecurityHub not enabled in region" + + # Verify that regions configuration was completely cleared + assert mock_integration.configuration["regions"] == {} + + # Verify save was called to persist the change + mock_integration.save.assert_called_once() + + # Verify RLS transaction was used correctly + # Should be called twice: once for getting provider info, once for resetting regions + assert mock_rls.call_count == 2 + mock_rls.assert_any_call(tenant_id) + + # Verify test_connection was called with integration credentials (not provider's) + mock_test_connection.assert_called_once_with( + aws_account_id="123456789012", + raise_on_exception=False, + aws_access_key_id="test_key_id", + aws_secret_access_key="test_secret_key", + ) + @patch("tasks.jobs.integrations.AwsProvider") @patch("tasks.jobs.integrations.SecurityHub.test_connection") @patch("tasks.jobs.integrations.initialize_prowler_provider")