Files
prowler/tests/providers/aws/services/awslambda/awslambda_service_test.py
Sebastian Nyberg 707584b2ef feat(aws): Add MFA flag if try to assume role in AWS (#2478)
Co-authored-by: Pepe Fagoaga <pepe@verica.io>
Co-authored-by: Sergio Garcia <sergargar1@gmail.com>
2023-06-13 17:18:10 +02:00

297 lines
11 KiB
Python

import io
import os
import tempfile
import zipfile
from re import search
from unittest.mock import patch
import mock
from boto3 import client, resource, session
from moto import mock_iam, mock_lambda, mock_s3
from moto.core import DEFAULT_ACCOUNT_ID
from prowler.providers.aws.lib.audit_info.models import AWS_Audit_Info
from prowler.providers.aws.services.awslambda.awslambda_service import AuthType, Lambda
from prowler.providers.common.models import Audit_Metadata
# Mock Test Region
AWS_REGION = "eu-west-1"
AWS_REGION_NORTH_VIRGINIA = "us-east-1"
def create_zip_file(code: str = "") -> io.BytesIO:
zip_output = io.BytesIO()
zip_file = zipfile.ZipFile(zip_output, "w", zipfile.ZIP_DEFLATED)
if not code:
zip_file.writestr(
"lambda_function.py",
"""
def lambda_handler(event, context):
print("custom log event")
return event
""",
)
else:
zip_file.writestr("lambda_function.py", code)
zip_file.close()
zip_output.seek(0)
return zip_output
def mock_request_get(_):
"""Mock requests.get() to get the Lambda Code in Zip Format"""
mock_resp = mock.MagicMock
mock_resp.status_code = 200
mock_resp.content = create_zip_file().read()
return mock_resp
# Mock generate_regional_clients()
def mock_generate_regional_clients(service, audit_info):
regional_client_eu_west_1 = audit_info.audit_session.client(
service, region_name=AWS_REGION
)
regional_client_us_east_1 = audit_info.audit_session.client(
service, region_name=AWS_REGION_NORTH_VIRGINIA
)
regional_client_eu_west_1.region = AWS_REGION
regional_client_us_east_1.region = AWS_REGION_NORTH_VIRGINIA
return {
AWS_REGION: regional_client_eu_west_1,
AWS_REGION_NORTH_VIRGINIA: regional_client_us_east_1,
}
@patch(
"prowler.providers.aws.services.awslambda.awslambda_service.generate_regional_clients",
new=mock_generate_regional_clients,
)
class Test_Lambda_Service:
def set_mocked_audit_info(self):
audit_info = AWS_Audit_Info(
session_config=None,
original_session=None,
audit_session=session.Session(
profile_name=None,
botocore_session=None,
),
audited_account=DEFAULT_ACCOUNT_ID,
audited_account_arn=f"arn:aws:iam::{DEFAULT_ACCOUNT_ID}:root",
audited_user_id=None,
audited_partition="aws",
audited_identity_arn=None,
profile=None,
profile_region=None,
credentials=None,
assumed_role_info=None,
audited_regions=None,
organizations_metadata=None,
audit_resources=None,
mfa_enabled=False,
audit_metadata=Audit_Metadata(
services_scanned=0,
# We need to set this check to call __list_functions__
expected_checks=["awslambda_function_no_secrets_in_code"],
completed_checks=0,
audit_progress=0,
),
)
return audit_info
# Test Lambda Client
def test__get_client__(self):
awslambda = Lambda(self.set_mocked_audit_info())
assert awslambda.regional_clients[AWS_REGION].__class__.__name__ == "Lambda"
# Test Lambda Session
def test__get_session__(self):
awslambda = Lambda(self.set_mocked_audit_info())
assert awslambda.session.__class__.__name__ == "Session"
# Test Lambda Service
def test__get_service__(self):
awslambda = Lambda(self.set_mocked_audit_info())
assert awslambda.service == "lambda"
@mock_lambda
@mock_iam
@mock_s3
def test__list_functions__(self):
# Create IAM Lambda Role
iam_client = client("iam", region_name=AWS_REGION)
iam_role = iam_client.create_role(
RoleName="test-lambda-role",
AssumeRolePolicyDocument="test-policy",
Path="/",
)["Role"]["Arn"]
# Create S3 Bucket
s3_client = resource("s3", region_name=AWS_REGION)
s3_client.create_bucket(
Bucket="test-bucket",
CreateBucketConfiguration={"LocationConstraint": AWS_REGION},
)
# Create Test Lambda 1
lambda_client = client("lambda", region_name=AWS_REGION)
lambda_name = "test-lambda"
resp = lambda_client.create_function(
FunctionName=lambda_name,
Runtime="python3.7",
Role=iam_role,
Handler="lambda_function.lambda_handler",
Code={"ZipFile": create_zip_file().read()},
Description="test lambda function",
Timeout=3,
MemorySize=128,
PackageType="ZIP",
Publish=True,
VpcConfig={
"SecurityGroupIds": ["sg-123abc"],
"SubnetIds": ["subnet-123abc"],
},
Environment={"Variables": {"db-password": "test-password"}},
Tags={"test": "test"},
)
lambda_arn_1 = resp["FunctionArn"]
# Update Lambda Policy
lambda_policy = {
"Version": "2012-10-17",
"Id": "default",
"Statement": [
{
"Action": "lambda:GetFunction",
"Principal": "*",
"Effect": "Allow",
"Resource": f"arn:aws:lambda:{AWS_REGION}:{DEFAULT_ACCOUNT_ID}:function:{lambda_name}",
"Sid": "test",
}
],
}
_ = lambda_client.add_permission(
FunctionName=lambda_name,
StatementId="test",
Action="lambda:GetFunction",
Principal="*",
)
# Create Function URL Config
_ = lambda_client.create_function_url_config(
FunctionName=lambda_name,
AuthType=AuthType.AWS_IAM.value,
Cors={
"AllowCredentials": True,
"AllowHeaders": [
"string",
],
"AllowMethods": [
"string",
],
"AllowOrigins": [
"*",
],
"ExposeHeaders": [
"string",
],
"MaxAge": 123,
},
)
# Create Test Lambda 2 (with the same attributes but different region)
lambda_client_2 = client("lambda", region_name=AWS_REGION_NORTH_VIRGINIA)
lambda_name = "test-lambda"
resp_2 = lambda_client_2.create_function(
FunctionName=lambda_name,
Runtime="python3.7",
Role=iam_role,
Handler="lambda_function.lambda_handler",
Code={"ZipFile": create_zip_file().read()},
Description="test lambda function",
Timeout=3,
MemorySize=128,
PackageType="ZIP",
Publish=True,
VpcConfig={
"SecurityGroupIds": ["sg-123abc"],
"SubnetIds": ["subnet-123abc"],
},
Environment={"Variables": {"db-password": "test-password"}},
Tags={"test": "test"},
)
lambda_arn_2 = resp_2["FunctionArn"]
with mock.patch(
"prowler.providers.aws.services.awslambda.awslambda_service.requests.get",
new=mock_request_get,
):
awslambda = Lambda(self.set_mocked_audit_info())
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].arn == lambda_arn_1
assert awslambda.functions[lambda_arn_1].runtime == "python3.7"
assert awslambda.functions[lambda_arn_1].environment == {
"db-password": "test-password"
}
assert awslambda.functions[lambda_arn_1].region == AWS_REGION
assert awslambda.functions[lambda_arn_1].policy == lambda_policy
assert awslambda.functions[lambda_arn_1].code
assert search(
f"s3://awslambda-{AWS_REGION}-tasks.s3-{AWS_REGION}.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
== AuthType.AWS_IAM
)
assert search(
"lambda-url.eu-west-1.on.aws",
awslambda.functions[lambda_arn_1].url_config.url,
)
assert awslambda.functions[lambda_arn_1].url_config.cors_config
assert awslambda.functions[
lambda_arn_1
].url_config.cors_config.allow_origins == ["*"]
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].arn == lambda_arn_2
assert awslambda.functions[lambda_arn_2].runtime == "python3.7"
assert awslambda.functions[lambda_arn_2].environment == {
"db-password": "test-password"
}
assert awslambda.functions[lambda_arn_2].region == AWS_REGION_NORTH_VIRGINIA
# Emtpy policy
assert awslambda.functions[lambda_arn_2].policy == {
"Id": "default",
"Statement": [],
"Version": "2012-10-17",
}
assert awslambda.functions[lambda_arn_2].code
assert search(
f"s3://awslambda-{AWS_REGION_NORTH_VIRGINIA}-tasks.s3-{AWS_REGION_NORTH_VIRGINIA}.amazonaws.com",
awslambda.functions[lambda_arn_2].code.location,
)