From 8b5c995486e0a59f97af62a6740867e42399471b Mon Sep 17 00:00:00 2001 From: Fennerr <41741346+Fennerr@users.noreply.github.com> Date: Wed, 13 Dec 2023 16:15:13 +0200 Subject: [PATCH] fix(lambda): memory leakage with lambda function code (#3167) Co-authored-by: Justin Moorcroft Co-authored-by: Pepe Fagoaga --- .../awslambda_function_no_secrets_in_code.py | 96 ++++++++------- .../services/awslambda/awslambda_service.py | 63 +++++----- ...lambda_function_no_secrets_in_code_test.py | 110 +++++++++--------- .../awslambda/awslambda_service_test.py | 79 ++++++------- 4 files changed, 174 insertions(+), 174 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 b90ba4e5..cc3d1075 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,57 +11,55 @@ 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: - report = Check_Report_AWS(self.metadata()) - report.region = function.region - report.resource_id = function.name - report.resource_arn = function.arn - report.resource_tags = function.tags + if awslambda_client.functions: + 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 + report.resource_arn = function.arn + report.resource_tags = function.tags - report.status = "PASS" - report.status_extended = ( - 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) - # List all files - files_in_zip = next(os.walk(tmp_dir_name))[2] - secrets_findings = [] - for file in files_in_zip: - secrets = SecretsCollection() - with default_settings(): - secrets.scan_file(f"{tmp_dir_name}/{file}") - detect_secrets_output = secrets.json() - if detect_secrets_output: - for ( - file_name - ) in ( - detect_secrets_output.keys() - ): # Appears that only 1 file is being scanned at a time, so could rework this - output_file_name = file_name.replace( - f"{tmp_dir_name}/", "" - ) - secrets_string = ", ".join( - [ - f"{secret['type']} on line {secret['line_number']}" - for secret in detect_secrets_output[file_name] - ] - ) - secrets_findings.append( - f"{output_file_name}: {secrets_string}" - ) + report.status = "PASS" + report.status_extended = ( + 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) + # List all files + files_in_zip = next(os.walk(tmp_dir_name))[2] + secrets_findings = [] + for file in files_in_zip: + secrets = SecretsCollection() + with default_settings(): + secrets.scan_file(f"{tmp_dir_name}/{file}") + detect_secrets_output = secrets.json() + if detect_secrets_output: + for ( + file_name + ) in ( + detect_secrets_output.keys() + ): # Appears that only 1 file is being scanned at a time, so could rework this + output_file_name = file_name.replace( + f"{tmp_dir_name}/", "" + ) + secrets_string = ", ".join( + [ + f"{secret['type']} on line {secret['line_number']}" + for secret in detect_secrets_output[ + file_name + ] + ] + ) + secrets_findings.append( + f"{output_file_name}: {secrets_string}" + ) - if secrets_findings: - final_output_string = "; ".join(secrets_findings) - report.status = "FAIL" - # report.status_extended = f"Potential {'secrets' if len(secrets_findings)>1 else 'secret'} found in Lambda function {function.name} code. {final_output_string}." - if len(secrets_findings) > 1: - report.status_extended = f"Potential secrets found in Lambda function {function.name} code -> {final_output_string}." - else: - report.status_extended = f"Potential secret found in Lambda function {function.name} code -> {final_output_string}." - # break // Don't break as there may be additional findings + if secrets_findings: + final_output_string = "; ".join(secrets_findings) + report.status = "FAIL" + report.status_extended = f"Potential {'secrets' if len(secrets_findings) > 1 else 'secret'} found in Lambda function {function.name} code -> {final_output_string}." - findings.append(report) + findings.append(report) return findings diff --git a/prowler/providers/aws/services/awslambda/awslambda_service.py b/prowler/providers/aws/services/awslambda/awslambda_service.py index 3b9d307e..72151973 100644 --- a/prowler/providers/aws/services/awslambda/awslambda_service.py +++ b/prowler/providers/aws/services/awslambda/awslambda_service.py @@ -1,6 +1,7 @@ import io import json import zipfile +from concurrent.futures import as_completed from enum import Enum from typing import Any, Optional @@ -21,15 +22,6 @@ class Lambda(AWSService): self.functions = {} self.__threading_call__(self.__list_functions__) self.__list_tags_for_resource__() - - # We only want to retrieve the Lambda code if the - # awslambda_function_no_secrets_in_code check is set - if ( - "awslambda_function_no_secrets_in_code" - in audit_info.audit_metadata.expected_checks - ): - self.__threading_call__(self.__get_function__) - self.__threading_call__(self.__get_policy__) self.__threading_call__(self.__get_function_url_config__) @@ -70,28 +62,45 @@ class Lambda(AWSService): f" {error}" ) - def __get_function__(self, regional_client): - logger.info("Lambda - Getting Function...") - try: - for function in self.functions.values(): - if function.region == regional_client.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 - self.functions[function.arn].code = LambdaCode( - location=code_location_uri, - code_zip=zipfile.ZipFile(io.BytesIO(raw_code_zip)), - ) + def __get_function_code__(self): + logger.info("Lambda - Getting Function Code...") + # Use a thread pool handle the queueing and execution of the __fetch_function_code__ tasks, up to max_workers tasks concurrently. + lambda_functions_to_fetch = { + self.thread_pool.submit( + self.__fetch_function_code__, function.name, function.region + ): function + for function in self.functions.values() + } + for fetched_lambda_code in as_completed(lambda_functions_to_fetch): + function = lambda_functions_to_fetch[fetched_lambda_code] + try: + function_code = fetched_lambda_code.result() + if function_code: + yield function, function_code + except Exception as error: + logger.error( + f"{function.region} -- {error.__class__.__name__}[{error.__traceback__.tb_lineno}]: {error}" + ) + + def __fetch_function_code__(self, function_name, function_region): + try: + 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 + return LambdaCode( + location=code_location_uri, + code_zip=zipfile.ZipFile(io.BytesIO(raw_code_zip)), + ) except Exception as error: logger.error( - f"{regional_client.region} --" - f" {error.__class__.__name__}[{error.__traceback__.tb_lineno}]:" - f" {error}" + f"{regional_client.region} -- {error.__class__.__name__}[{error.__traceback__.tb_lineno}]: {error}" ) + raise def __get_policy__(self, regional_client): logger.info("Lambda - Getting Policy...") diff --git a/tests/providers/aws/services/awslambda/awslambda_function_no_secrets_in_code/awslambda_function_no_secrets_in_code_test.py b/tests/providers/aws/services/awslambda/awslambda_function_no_secrets_in_code/awslambda_function_no_secrets_in_code_test.py index 223828cd..fa0ec88c 100644 --- a/tests/providers/aws/services/awslambda/awslambda_function_no_secrets_in_code/awslambda_function_no_secrets_in_code_test.py +++ b/tests/providers/aws/services/awslambda/awslambda_function_no_secrets_in_code/awslambda_function_no_secrets_in_code_test.py @@ -1,8 +1,6 @@ import zipfile from unittest import mock -from awslambda_service_test import create_zip_file - from prowler.providers.aws.services.awslambda.awslambda_service import ( Function, LambdaCode, @@ -12,6 +10,53 @@ from tests.providers.aws.audit_info_utils import ( AWS_REGION_US_EAST_1, set_mocked_aws_audit_info, ) +from tests.providers.aws.services.awslambda.awslambda_service_test import ( + create_zip_file, +) + +LAMBDA_FUNCTION_NAME = "test-lambda" +LAMBDA_FUNCTION_RUNTIME = "nodejs4.3" +LAMBDA_FUNCTION_ARN = f"arn:aws:lambda:{AWS_REGION_US_EAST_1}:{AWS_ACCOUNT_NUMBER}:function/{LAMBDA_FUNCTION_NAME}" +LAMBDA_FUNCTION_CODE_WITH_SECRETS = """ +def lambda_handler(event, context): + db_password = "test-password" + print("custom log event") + return event +""" +LAMBDA_FUNCTION_CODE_WITHOUT_SECRETS = """ +def lambda_handler(event, context): + print("custom log event") + return event +""" + + +def create_lambda_function() -> Function: + return Function( + name=LAMBDA_FUNCTION_NAME, + security_groups=[], + arn=LAMBDA_FUNCTION_ARN, + region=AWS_REGION_US_EAST_1, + runtime=LAMBDA_FUNCTION_RUNTIME, + ) + + +def get_lambda_code_with_secrets(code): + return LambdaCode( + location="", + code_zip=zipfile.ZipFile(create_zip_file(code)), + ) + + +def mock__get_function_code__with_secrets(): + yield create_lambda_function(), get_lambda_code_with_secrets( + LAMBDA_FUNCTION_CODE_WITH_SECRETS + ) + + +def mock__get_function_code__without_secrets(): + yield create_lambda_function(), get_lambda_code_with_secrets( + LAMBDA_FUNCTION_CODE_WITHOUT_SECRETS + ) class Test_awslambda_function_no_secrets_in_code: @@ -38,29 +83,8 @@ class Test_awslambda_function_no_secrets_in_code: def test_function_code_with_secrets(self): lambda_client = mock.MagicMock - function_name = "test-lambda" - function_runtime = "nodejs4.3" - function_arn = f"arn:aws:lambda:{AWS_REGION_US_EAST_1}:{AWS_ACCOUNT_NUMBER}:function/{function_name}" - code_with_secrets = """ - def lambda_handler(event, context): - db_password = "test-password" - print("custom log event") - return event - """ - lambda_client.functions = { - "function_name": Function( - name=function_name, - security_groups=[], - arn=function_arn, - region=AWS_REGION_US_EAST_1, - runtime=function_runtime, - code=LambdaCode( - location="", - code_zip=zipfile.ZipFile(create_zip_file(code_with_secrets)), - ), - ) - } - + lambda_client.functions = {LAMBDA_FUNCTION_ARN: create_lambda_function()} + lambda_client.__get_function_code__ = mock__get_function_code__with_secrets with mock.patch( "prowler.providers.aws.lib.audit_info.audit_info.current_audit_info", set_mocked_aws_audit_info(), @@ -78,38 +102,20 @@ class Test_awslambda_function_no_secrets_in_code: assert len(result) == 1 assert result[0].region == AWS_REGION_US_EAST_1 - assert result[0].resource_id == function_name - assert result[0].resource_arn == function_arn + assert result[0].resource_id == LAMBDA_FUNCTION_NAME + assert result[0].resource_arn == LAMBDA_FUNCTION_ARN assert result[0].status == "FAIL" assert ( result[0].status_extended - == f"Potential secret found in Lambda function {function_name} code -> lambda_function.py: Secret Keyword on line 3." + == f"Potential secret found in Lambda function {LAMBDA_FUNCTION_NAME} code -> lambda_function.py: Secret Keyword on line 3." ) assert result[0].resource_tags == [] def test_function_code_without_secrets(self): lambda_client = mock.MagicMock - function_name = "test-lambda" - function_runtime = "nodejs4.3" - function_arn = f"arn:aws:lambda:{AWS_REGION_US_EAST_1}:{AWS_ACCOUNT_NUMBER}:function/{function_name}" - code_with_secrets = """ - def lambda_handler(event, context): - print("custom log event") - return event - """ - lambda_client.functions = { - "function_name": Function( - name=function_name, - security_groups=[], - arn=function_arn, - region=AWS_REGION_US_EAST_1, - runtime=function_runtime, - code=LambdaCode( - location="", - code_zip=zipfile.ZipFile(create_zip_file(code_with_secrets)), - ), - ) - } + lambda_client.functions = {LAMBDA_FUNCTION_ARN: create_lambda_function()} + + lambda_client.__get_function_code__ = mock__get_function_code__without_secrets with mock.patch( "prowler.providers.aws.lib.audit_info.audit_info.current_audit_info", @@ -128,11 +134,11 @@ class Test_awslambda_function_no_secrets_in_code: assert len(result) == 1 assert result[0].region == AWS_REGION_US_EAST_1 - assert result[0].resource_id == function_name - assert result[0].resource_arn == function_arn + assert result[0].resource_id == LAMBDA_FUNCTION_NAME + assert result[0].resource_arn == LAMBDA_FUNCTION_ARN assert result[0].status == "PASS" assert ( result[0].status_extended - == f"No secrets found in Lambda function {function_name} code." + == f"No secrets found in Lambda function {LAMBDA_FUNCTION_NAME} code." ) assert result[0].resource_tags == [] diff --git a/tests/providers/aws/services/awslambda/awslambda_service_test.py b/tests/providers/aws/services/awslambda/awslambda_service_test.py index 6933d1bd..5a4ff328 100644 --- a/tests/providers/aws/services/awslambda/awslambda_service_test.py +++ b/tests/providers/aws/services/awslambda/awslambda_service_test.py @@ -17,6 +17,11 @@ from tests.providers.aws.audit_info_utils import ( set_mocked_aws_audit_info, ) +LAMBDA_FUNCTION_CODE = """def lambda_handler(event, context): +print("custom log event") +return event + """ + def create_zip_file(code: str = "") -> io.BytesIO: zip_output = io.BytesIO() @@ -24,11 +29,7 @@ def create_zip_file(code: str = "") -> io.BytesIO: if not code: zip_file.writestr( "lambda_function.py", - """ - def lambda_handler(event, context): - print("custom log event") - return event - """, + LAMBDA_FUNCTION_CODE, ) else: zip_file.writestr("lambda_function.py", code) @@ -103,9 +104,9 @@ class Test_Lambda_Service: ) # Create Test Lambda 1 lambda_client = client("lambda", region_name=AWS_REGION_EU_WEST_1) - lambda_name = "test-lambda" + lambda_name_1 = "test-lambda-1" resp = lambda_client.create_function( - FunctionName=lambda_name, + FunctionName=lambda_name_1, Runtime="python3.7", Role=iam_role, Handler="lambda_function.lambda_handler", @@ -132,20 +133,20 @@ class Test_Lambda_Service: "Action": "lambda:GetFunction", "Principal": "*", "Effect": "Allow", - "Resource": f"arn:aws:lambda:{AWS_REGION_EU_WEST_1}:{AWS_ACCOUNT_NUMBER}:function:{lambda_name}", + "Resource": f"arn:aws:lambda:{AWS_REGION_EU_WEST_1}:{AWS_ACCOUNT_NUMBER}:function:{lambda_name_1}", "Sid": "test", } ], } _ = lambda_client.add_permission( - FunctionName=lambda_name, + FunctionName=lambda_name_1, StatementId="test", Action="lambda:GetFunction", Principal="*", ) # Create Function URL Config _ = lambda_client.create_function_url_config( - FunctionName=lambda_name, + FunctionName=lambda_name_1, AuthType=AuthType.AWS_IAM.value, Cors={ "AllowCredentials": True, @@ -167,9 +168,9 @@ class Test_Lambda_Service: # Create Test Lambda 2 (with the same attributes but different region) lambda_client_2 = client("lambda", region_name=AWS_REGION_US_EAST_1) - lambda_name = "test-lambda" + lambda_name_2 = "test-lambda-2" resp_2 = lambda_client_2.create_function( - FunctionName=lambda_name, + FunctionName=lambda_name_2, Runtime="python3.7", Role=iam_role, Handler="lambda_function.lambda_handler", @@ -193,15 +194,12 @@ class Test_Lambda_Service: new=mock_request_get, ): awslambda = Lambda( - set_mocked_aws_audit_info( - audited_regions=[AWS_REGION_US_EAST_1], - expected_checks=["awslambda_function_no_secrets_in_code"], - ) + set_mocked_aws_audit_info(audited_regions=[AWS_REGION_US_EAST_1]) ) assert awslambda.functions assert len(awslambda.functions) == 2 # Lambda 1 - assert awslambda.functions[lambda_arn_1].name == lambda_name + assert awslambda.functions[lambda_arn_1].name == lambda_name_1 assert awslambda.functions[lambda_arn_1].arn == lambda_arn_1 assert awslambda.functions[lambda_arn_1].runtime == "python3.7" assert awslambda.functions[lambda_arn_1].environment == { @@ -210,12 +208,6 @@ class Test_Lambda_Service: assert awslambda.functions[lambda_arn_1].region == AWS_REGION_EU_WEST_1 assert awslambda.functions[lambda_arn_1].policy == lambda_policy - assert awslambda.functions[lambda_arn_1].code - assert search( - f"s3://awslambda-{AWS_REGION_EU_WEST_1}-tasks.s3-{AWS_REGION_EU_WEST_1}.amazonaws.com", - awslambda.functions[lambda_arn_1].code.location, - ) - assert awslambda.functions[lambda_arn_1].url_config assert ( awslambda.functions[lambda_arn_1].url_config.auth_type @@ -233,25 +225,8 @@ class Test_Lambda_Service: assert awslambda.functions[lambda_arn_1].tags == [{"test": "test"}] - # Pending ZipFile tests - with tempfile.TemporaryDirectory() as tmp_dir_name: - awslambda.functions[lambda_arn_1].code.code_zip.extractall(tmp_dir_name) - files_in_zip = next(os.walk(tmp_dir_name))[2] - assert len(files_in_zip) == 1 - assert files_in_zip[0] == "lambda_function.py" - with open(f"{tmp_dir_name}/{files_in_zip[0]}", "r") as lambda_code_file: - _ = lambda_code_file - # assert ( - # lambda_code_file.read() - # == """ - # def lambda_handler(event, context): - # print("custom log event") - # return event - # """ - # ) - # Lambda 2 - assert awslambda.functions[lambda_arn_2].name == lambda_name + assert awslambda.functions[lambda_arn_2].name == lambda_name_2 assert awslambda.functions[lambda_arn_2].arn == lambda_arn_2 assert awslambda.functions[lambda_arn_2].runtime == "python3.7" assert awslambda.functions[lambda_arn_2].environment == { @@ -265,8 +240,20 @@ class Test_Lambda_Service: "Version": "2012-10-17", } - assert awslambda.functions[lambda_arn_2].code - assert search( - f"s3://awslambda-{AWS_REGION_US_EAST_1}-tasks.s3-{AWS_REGION_US_EAST_1}.amazonaws.com", - awslambda.functions[lambda_arn_2].code.location, - ) + # Lambda Code + with tempfile.TemporaryDirectory() as tmp_dir_name: + for function, function_code in awslambda.__get_function_code__(): + if function.arn == lambda_arn_1 or function.arn == lambda_arn_2: + assert search( + f"s3://awslambda-{function.region}-tasks.s3-{function.region}.amazonaws.com", + function_code.location, + ) + assert function_code + function_code.code_zip.extractall(tmp_dir_name) + files_in_zip = next(os.walk(tmp_dir_name))[2] + assert len(files_in_zip) == 1 + assert files_in_zip[0] == "lambda_function.py" + with open( + f"{tmp_dir_name}/{files_in_zip[0]}", "r" + ) as lambda_code_file: + assert lambda_code_file.read() == LAMBDA_FUNCTION_CODE