96 lines
2.9 KiB
Python
96 lines
2.9 KiB
Python
|
import io
|
||
|
|
import json
|
||
|
|
import zipfile
|
||
|
|
from unittest import SkipTest
|
||
|
|
from unittest.mock import patch
|
||
|
|
|
||
|
|
import boto3
|
||
|
|
from botocore.exceptions import ClientError
|
||
|
|
|
||
|
|
from moto import mock_aws, settings
|
||
|
|
|
||
|
|
secret_steps = []
|
||
|
|
|
||
|
|
|
||
|
|
def mock_lambda_invoke(*args, **kwarg):
|
||
|
|
secret_steps.append(json.loads(kwarg["body"])["Step"])
|
||
|
|
return "n/a"
|
||
|
|
|
||
|
|
|
||
|
|
@mock_aws(config={"lambda": {"use_docker": False}})
|
||
|
|
@patch(
|
||
|
|
"moto.awslambda_simple.models.LambdaSimpleBackend.invoke", new=mock_lambda_invoke
|
||
|
|
)
|
||
|
|
def test_simple_lambda_is_invoked():
|
||
|
|
if not settings.TEST_DECORATOR_MODE:
|
||
|
|
raise SkipTest("Can only test patched code in DecoratorMode")
|
||
|
|
sm_client = boto3.client("secretsmanager", region_name="us-east-1")
|
||
|
|
secret_arn = sm_client.create_secret(Name="some", SecretString="secret")["ARN"]
|
||
|
|
|
||
|
|
lambda_res = create_mock_rotator_lambda()
|
||
|
|
sm_client.rotate_secret(
|
||
|
|
SecretId=secret_arn,
|
||
|
|
RotationLambdaARN=lambda_res["FunctionArn"],
|
||
|
|
RotationRules={"AutomaticallyAfterDays": 1, "Duration": "1h"},
|
||
|
|
RotateImmediately=True,
|
||
|
|
)
|
||
|
|
assert secret_steps == ["createSecret", "setSecret", "testSecret", "finishSecret"]
|
||
|
|
secret_steps.clear()
|
||
|
|
|
||
|
|
|
||
|
|
@mock_aws(config={"lambda": {"use_docker": False}})
|
||
|
|
@patch(
|
||
|
|
"moto.awslambda_simple.models.LambdaSimpleBackend.invoke", new=mock_lambda_invoke
|
||
|
|
)
|
||
|
|
def test_simple_lambda_is_invoked__do_not_rotate_immediately():
|
||
|
|
if not settings.TEST_DECORATOR_MODE:
|
||
|
|
raise SkipTest("Can only test patched code in DecoratorMode")
|
||
|
|
sm_client = boto3.client("secretsmanager", region_name="us-east-1")
|
||
|
|
secret_arn = sm_client.create_secret(Name="some", SecretString="secret")["ARN"]
|
||
|
|
|
||
|
|
lambda_res = create_mock_rotator_lambda()
|
||
|
|
sm_client.rotate_secret(
|
||
|
|
SecretId=secret_arn,
|
||
|
|
RotationLambdaARN=lambda_res["FunctionArn"],
|
||
|
|
RotationRules={"AutomaticallyAfterDays": 1, "Duration": "1h"},
|
||
|
|
RotateImmediately=False,
|
||
|
|
)
|
||
|
|
assert secret_steps == ["testSecret"]
|
||
|
|
secret_steps.clear()
|
||
|
|
|
||
|
|
|
||
|
|
def mock_lambda_zip():
|
||
|
|
code = """
|
||
|
|
def lambda_handler(event, context):
|
||
|
|
return event
|
||
|
|
"""
|
||
|
|
zip_output = io.BytesIO()
|
||
|
|
zip_file = zipfile.ZipFile(zip_output, "w", zipfile.ZIP_DEFLATED)
|
||
|
|
zip_file.writestr("lambda_function.py", code)
|
||
|
|
zip_file.close()
|
||
|
|
zip_output.seek(0)
|
||
|
|
return zip_output.read()
|
||
|
|
|
||
|
|
|
||
|
|
def create_mock_rotator_lambda():
|
||
|
|
client = boto3.client("lambda", region_name="us-east-1")
|
||
|
|
return client.create_function(
|
||
|
|
FunctionName="mock-rotator",
|
||
|
|
Runtime="python3.9",
|
||
|
|
Role=get_mock_role_arn(),
|
||
|
|
Handler="lambda_function.lambda_handler",
|
||
|
|
Code={"ZipFile": mock_lambda_zip()},
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
def get_mock_role_arn():
|
||
|
|
iam = boto3.client("iam", region_name="us-east-1")
|
||
|
|
try:
|
||
|
|
return iam.get_role(RoleName="my-role")["Role"]["Arn"]
|
||
|
|
except ClientError:
|
||
|
|
return iam.create_role(
|
||
|
|
RoleName="my-role",
|
||
|
|
AssumeRolePolicyDocument="some policy",
|
||
|
|
Path="/my-path/",
|
||
|
|
)["Role"]["Arn"]
|