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

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

View File

@@ -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 == []

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