From 240f57278af236683a80677e37a65d1f0dfda235 Mon Sep 17 00:00:00 2001
From: Prowler Bot <bot@prowler.com>
Date: Fri, 17 Jan 2025 15:27:01 +0100
Subject: [PATCH] fix(apigatewayv2): managed exception `NotFoundException`
 (#6589)

Co-authored-by: Hugo Pereira Brito <101209179+HugoPBrito@users.noreply.github.com>
---
 .../apigatewayv2/apigatewayv2_service.py      | 11 ++++++-
 .../apigatewayv2/apigatewayv2_service_test.py | 31 +++++++++++++++++++
 2 files changed, 41 insertions(+), 1 deletion(-)

diff --git a/prowler/providers/aws/services/apigatewayv2/apigatewayv2_service.py b/prowler/providers/aws/services/apigatewayv2/apigatewayv2_service.py
index 0ee3d812528..cb8d88e5ef8 100644
--- a/prowler/providers/aws/services/apigatewayv2/apigatewayv2_service.py
+++ b/prowler/providers/aws/services/apigatewayv2/apigatewayv2_service.py
@@ -1,5 +1,6 @@
 from typing import Optional
 
+from botocore.exceptions import ClientError
 from pydantic import BaseModel
 
 from prowler.lib.logger import logger
@@ -7,7 +8,6 @@
 from prowler.providers.aws.lib.service.service import AWSService
 
 
-################## ApiGatewayV2
 class ApiGatewayV2(AWSService):
     def __init__(self, provider):
         # Call AWSService's __init__
@@ -71,6 +71,15 @@ def _get_stages(self):
                             tags=[stage.get("Tags")],
                         )
                     )
+        except ClientError as error:
+            if error.response["Error"]["Code"] == "NotFoundException":
+                logger.warning(
+                    f"{regional_client.region} -- {error.__class__.__name__}[{error.__traceback__.tb_lineno}]: {error}"
+                )
+            else:
+                logger.error(
+                    f"{regional_client.region} -- {error.__class__.__name__}[{error.__traceback__.tb_lineno}]: {error}"
+                )
         except Exception as error:
             logger.error(
                 f"{error.__class__.__name__}:{error.__traceback__.tb_lineno} -- {error}"
diff --git a/tests/providers/aws/services/apigatewayv2/apigatewayv2_service_test.py b/tests/providers/aws/services/apigatewayv2/apigatewayv2_service_test.py
index 199a257b1d7..77381d3cb4a 100644
--- a/tests/providers/aws/services/apigatewayv2/apigatewayv2_service_test.py
+++ b/tests/providers/aws/services/apigatewayv2/apigatewayv2_service_test.py
@@ -25,6 +25,16 @@ def mock_make_api_call(self, operation_name, kwarg):
     if operation_name == "GetAuthorizers":
         return {"Items": [{"AuthorizerId": "authorizer-id", "Name": "test-authorizer"}]}
     elif operation_name == "GetStages":
+        if kwarg["ApiId"] == "not-found-api":
+            raise botocore.exceptions.ClientError(
+                {
+                    "Error": {
+                        "Code": "NotFoundException",
+                        "Message": "API not found",
+                    }
+                },
+                "GetStages",
+            )
         return {
             "Items": [
                 {
@@ -120,3 +130,24 @@ def test_get_stages(self):
         aws_provider = set_mocked_aws_provider([AWS_REGION_US_EAST_1])
         apigatewayv2 = ApiGatewayV2(aws_provider)
         assert apigatewayv2.apis[0].stages[0].logging is True
+
+    # Test ApiGatewayV2 Get Stages with NotFoundException
+    @mock_aws
+    @patch("prowler.providers.aws.services.apigatewayv2.apigatewayv2_service.logger")
+    def test_get_stages_not_found_exception(self, mock_logger):
+        # Generate ApiGatewayV2 Client
+        apigatewayv2_client = client("apigatewayv2", region_name=AWS_REGION_US_EAST_1)
+        # Create ApiGatewayV2 Rest API
+        apigatewayv2_client.create_api(Name="test-api", ProtocolType="HTTP")
+
+        aws_provider = set_mocked_aws_provider([AWS_REGION_US_EAST_1])
+        apigatewayv2 = ApiGatewayV2(aws_provider)
+
+        # Force API ID to trigger NotFoundException
+        apigatewayv2.apis[0].id = "not-found-api"
+
+        # Call _get_stages to trigger the exception
+        apigatewayv2._get_stages()
+
+        mock_logger.warning.assert_called_once()
+        assert "NotFoundException" in mock_logger.warning.call_args[0][0]