From dedb03cc6e103366651e7decb36140dc238746eb Mon Sep 17 00:00:00 2001
From: Justin Moorcroft <justin.moorcroft@mwrcybersec.com>
Date: Tue, 5 Dec 2023 22:11:32 +0200
Subject: [PATCH] initial fix for issue #2516

---
 .../awslambda_function_no_secrets_in_code.py  |  6 ++---
 .../services/awslambda/awslambda_service.py   | 27 ++++++++++---------
 2 files changed, 17 insertions(+), 16 deletions(-)

diff --git a/prowler/providers/aws/services/awslambda/awslambda_function_no_secrets_in_code/awslambda_function_no_secrets_in_code.py b/prowler/providers/aws/services/awslambda/awslambda_function_no_secrets_in_code/awslambda_function_no_secrets_in_code.py
index b90ba4e5e3..5034cc76b5 100644
--- a/prowler/providers/aws/services/awslambda/awslambda_function_no_secrets_in_code/awslambda_function_no_secrets_in_code.py
+++ b/prowler/providers/aws/services/awslambda/awslambda_function_no_secrets_in_code/awslambda_function_no_secrets_in_code.py
@@ -11,8 +11,8 @@ from prowler.providers.aws.services.awslambda.awslambda_client import awslambda_
 class awslambda_function_no_secrets_in_code(Check):
     def execute(self):
         findings = []
-        for function in awslambda_client.functions.values():
-            if function.code:
+        for function, function_code in awslambda_client.__get_function_code__():
+            if function_code:
                 report = Check_Report_AWS(self.metadata())
                 report.region = function.region
                 report.resource_id = function.name
@@ -24,7 +24,7 @@ class awslambda_function_no_secrets_in_code(Check):
                     f"No secrets found in Lambda function {function.name} code."
                 )
                 with tempfile.TemporaryDirectory() as tmp_dir_name:
-                    function.code.code_zip.extractall(tmp_dir_name)
+                    function_code.code_zip.extractall(tmp_dir_name)
                     # List all files
                     files_in_zip = next(os.walk(tmp_dir_name))[2]
                     secrets_findings = []
diff --git a/prowler/providers/aws/services/awslambda/awslambda_service.py b/prowler/providers/aws/services/awslambda/awslambda_service.py
index 3b9d307e1c..4cff2e2254 100644
--- a/prowler/providers/aws/services/awslambda/awslambda_service.py
+++ b/prowler/providers/aws/services/awslambda/awslambda_service.py
@@ -70,21 +70,22 @@ class Lambda(AWSService):
                 f" {error}"
             )
 
-    def __get_function__(self, regional_client):
-        logger.info("Lambda - Getting Function...")
+    def __get_function_code__(self, function):
+        logger.info("Lambda - Getting Function Code...")
         try:
             for function in self.functions.values():
-                if function.region == regional_client.region:
-                    function_information = regional_client.get_function(
-                        FunctionName=function.name
+                regional_client = self.regional_clients[function.region]
+                function_information = regional_client.get_function(
+                    FunctionName=function.name
+                )
+                if "Location" in function_information["Code"]:
+                    code_location_uri = function_information["Code"]["Location"]
+                    raw_code_zip = requests.get(code_location_uri).content
+                    function_code = LambdaCode(
+                        location=code_location_uri,
+                        code_zip=zipfile.ZipFile(io.BytesIO(raw_code_zip)),
                     )
-                    if "Location" in function_information["Code"]:
-                        code_location_uri = function_information["Code"]["Location"]
-                        raw_code_zip = requests.get(code_location_uri).content
-                        self.functions[function.arn].code = LambdaCode(
-                            location=code_location_uri,
-                            code_zip=zipfile.ZipFile(io.BytesIO(raw_code_zip)),
-                        )
+                    yield function, function_code
 
         except Exception as error:
             logger.error(
@@ -92,7 +93,7 @@ class Lambda(AWSService):
                 f" {error.__class__.__name__}[{error.__traceback__.tb_lineno}]:"
                 f" {error}"
             )
-
+            
     def __get_policy__(self, regional_client):
         logger.info("Lambda - Getting Policy...")
         try: