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

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

View File

@@ -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...")

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