ghndrx/prowler
1
0
Fork 0
mirror of https://github.com/ghndrx/prowler.git synced 2026-02-10 23:05:05 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix(lambda): memory leakage with lambda function code (#3167)

Browse Source 
Co-authored-by: Justin Moorcroft <justin.moorcroft@mwrcybersec.com>
Co-authored-by: Pepe Fagoaga <pepe@verica.io>
This commit is contained in:
Fennerr
2023-12-13 16:15:13 +02:00
committed by GitHub
parent 4410f2a582
commit 8b5c995486
4 changed files with 174 additions and 174 deletions
Download Patch File Download Diff File Expand all files Collapse all files

79
tests/providers/aws/services/awslambda/awslambda_service_test.py
View File

@@ -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