From fc9c2509228c0bd04300347d50d46d0ff1c7087f Mon Sep 17 00:00:00 2001 From: Chris Keogh Date: Tue, 3 Oct 2017 13:33:50 +1300 Subject: [PATCH 1/3] add basic awslambda get_policy --- moto/awslambda/responses.py | 14 ++++++++++++++ moto/awslambda/urls.py | 3 ++- tests/test_awslambda/test_lambda.py | 23 +++++++++++++++++++++++ 3 files changed, 39 insertions(+), 1 deletion(-) diff --git a/moto/awslambda/responses.py b/moto/awslambda/responses.py index cf92e66f4..972cd7a4e 100644 --- a/moto/awslambda/responses.py +++ b/moto/awslambda/responses.py @@ -57,6 +57,20 @@ class LambdaResponse(BaseResponse): else: raise ValueError("Cannot handle {0} request".format(request.method)) + def policy(self, request, full_url, headers): + if request.method == 'GET': + return self._get_policy(request, full_url, headers) + + def _get_policy(self, request, full_url, headers): + lambda_backend = self.get_lambda_backend(full_url) + + path = request.path if hasattr(request, 'path') else request.path_url + function_name = path.split('/')[-2] + if lambda_backend.has_function(function_name): + return 200, {}, json.dumps(dict(Policy='test_policy')) + else: + return 404, {}, "{}" + def _invoke(self, request, full_url): response_headers = {} lambda_backend = self.get_lambda_backend(full_url) diff --git a/moto/awslambda/urls.py b/moto/awslambda/urls.py index 0fec24bab..005785f19 100644 --- a/moto/awslambda/urls.py +++ b/moto/awslambda/urls.py @@ -12,5 +12,6 @@ url_paths = { r'{0}/(?P[^/]+)/functions/(?P[\w_-]+)/?$': response.function, r'{0}/(?P[^/]+)/functions/(?P[\w_-]+)/invocations/?$': response.invoke, r'{0}/(?P[^/]+)/functions/(?P[\w_-]+)/invoke-async/?$': response.invoke_async, - r'{0}/(?P[^/]+)/tags/(?P.+)': response.tag + r'{0}/(?P[^/]+)/tags/(?P.+)': response.tag, + r'{0}/(?P[^/]+)/functions/(?P[\w_-]+)/policy/?$': response.policy } diff --git a/tests/test_awslambda/test_lambda.py b/tests/test_awslambda/test_lambda.py index 6b67ce0f0..c7fee2745 100644 --- a/tests/test_awslambda/test_lambda.py +++ b/tests/test_awslambda/test_lambda.py @@ -645,3 +645,26 @@ def test_get_function_created_with_zipfile(): } }, ) + +@mock_lambda +def get_function_policy(): + conn = boto3.client('lambda', 'us-west-2') + zip_content = get_test_zip_file1() + result = conn.create_function( + FunctionName='testFunction', + Runtime='python2.7', + Role='test-iam-role', + Handler='lambda_function.handler', + Code={ + 'ZipFile': zip_content, + }, + Description='test lambda function', + Timeout=3, + MemorySize=128, + Publish=True, + ) + + response = conn.get_policy( + FunctionName='testFunction' + ) + assert response['Policy'] == 'test_policy' \ No newline at end of file From b994cf5291888858c46609449ff0d5d940ea5a59 Mon Sep 17 00:00:00 2001 From: Chris Keogh Date: Tue, 3 Oct 2017 13:54:37 +1300 Subject: [PATCH 2/3] add more realistic policy --- moto/awslambda/responses.py | 5 ++++- tests/test_awslambda/test_lambda.py | 6 +++++- 2 files changed, 9 insertions(+), 2 deletions(-) diff --git a/moto/awslambda/responses.py b/moto/awslambda/responses.py index 972cd7a4e..94c381f5e 100644 --- a/moto/awslambda/responses.py +++ b/moto/awslambda/responses.py @@ -67,7 +67,10 @@ class LambdaResponse(BaseResponse): path = request.path if hasattr(request, 'path') else request.path_url function_name = path.split('/')[-2] if lambda_backend.has_function(function_name): - return 200, {}, json.dumps(dict(Policy='test_policy')) + policy = ("{\"Statement\":[{\"Action\":[\"lambda:InvokeFunction\"]," + "\"Resource\":\"arn:aws:lambda:us-west-2:account-id:function:helloworld\"," + "\"Effect\":\"Allow\",\"Principal\":{\"AWS\":\"account-id\"},\"Sid\":\"3\"}]}") + return 200, {}, json.dumps(dict(Policy=policy)) else: return 404, {}, "{}" diff --git a/tests/test_awslambda/test_lambda.py b/tests/test_awslambda/test_lambda.py index c7fee2745..163fa306f 100644 --- a/tests/test_awslambda/test_lambda.py +++ b/tests/test_awslambda/test_lambda.py @@ -667,4 +667,8 @@ def get_function_policy(): response = conn.get_policy( FunctionName='testFunction' ) - assert response['Policy'] == 'test_policy' \ No newline at end of file + + assert 'Policy' in response + assert isinstance(response['Policy'], str) + res = json.loads(response['Policy']) + assert res['Statement'][0]['Action'] == ['lambda:InvokeFunction'] \ No newline at end of file From 9bb07e6b6e42811fa768849618a7c32d9f55d4d1 Mon Sep 17 00:00:00 2001 From: Chris Keogh Date: Tue, 3 Oct 2017 15:23:00 +1300 Subject: [PATCH 3/3] add awslambda.add_permission --- moto/awslambda/models.py | 4 +++ moto/awslambda/responses.py | 20 ++++++++++--- tests/test_awslambda/test_lambda.py | 46 ++++++++++++++++++++++++++++- 3 files changed, 65 insertions(+), 5 deletions(-) diff --git a/moto/awslambda/models.py b/moto/awslambda/models.py index d22d1a7f4..935abbcd6 100644 --- a/moto/awslambda/models.py +++ b/moto/awslambda/models.py @@ -132,6 +132,7 @@ class LambdaFunction(BaseModel): self.logs_backend = logs_backends[self.region] self.environment_vars = spec.get('Environment', {}).get('Variables', {}) self.docker_client = docker.from_env() + self.policy = "" # Unfortunately mocking replaces this method w/o fallback enabled, so we # need to replace it if we detect it's been mocked @@ -527,6 +528,9 @@ class LambdaBackend(BaseBackend): pass # Don't care + def add_policy(self, function_name, policy): + self.get_function(function_name).policy = policy + def do_validate_s3(): return os.environ.get('VALIDATE_LAMBDA_S3', '') in ['', '1', 'true'] diff --git a/moto/awslambda/responses.py b/moto/awslambda/responses.py index 94c381f5e..5215f63c5 100644 --- a/moto/awslambda/responses.py +++ b/moto/awslambda/responses.py @@ -60,6 +60,20 @@ class LambdaResponse(BaseResponse): def policy(self, request, full_url, headers): if request.method == 'GET': return self._get_policy(request, full_url, headers) + if request.method == 'POST': + return self._add_policy(request, full_url, headers) + + def _add_policy(self, request, full_url, headers): + lambda_backend = self.get_lambda_backend(full_url) + + path = request.path if hasattr(request, 'path') else request.path_url + function_name = path.split('/')[-2] + if lambda_backend.has_function(function_name): + policy = request.body.decode('utf8') + lambda_backend.add_policy(function_name, policy) + return 200, {}, json.dumps(dict(Statement=policy)) + else: + return 404, {}, "{}" def _get_policy(self, request, full_url, headers): lambda_backend = self.get_lambda_backend(full_url) @@ -67,10 +81,8 @@ class LambdaResponse(BaseResponse): path = request.path if hasattr(request, 'path') else request.path_url function_name = path.split('/')[-2] if lambda_backend.has_function(function_name): - policy = ("{\"Statement\":[{\"Action\":[\"lambda:InvokeFunction\"]," - "\"Resource\":\"arn:aws:lambda:us-west-2:account-id:function:helloworld\"," - "\"Effect\":\"Allow\",\"Principal\":{\"AWS\":\"account-id\"},\"Sid\":\"3\"}]}") - return 200, {}, json.dumps(dict(Policy=policy)) + function = lambda_backend.get_function(function_name) + return 200, {}, json.dumps(dict(Policy="{\"Statement\":[" + function.policy + "]}")) else: return 404, {}, "{}" diff --git a/tests/test_awslambda/test_lambda.py b/tests/test_awslambda/test_lambda.py index 163fa306f..317e9f4a2 100644 --- a/tests/test_awslambda/test_lambda.py +++ b/tests/test_awslambda/test_lambda.py @@ -646,6 +646,39 @@ def test_get_function_created_with_zipfile(): }, ) +@mock_lambda +def add_function_permission(): + conn = boto3.client('lambda', 'us-west-2') + zip_content = get_test_zip_file1() + result = conn.create_function( + FunctionName='testFunction', + Runtime='python2.7', + Role='test-iam-role', + Handler='lambda_function.handler', + Code={ + 'ZipFile': zip_content, + }, + Description='test lambda function', + Timeout=3, + MemorySize=128, + Publish=True, + ) + + response = conn.add_permission( + FunctionName='testFunction', + StatementId='1', + Action="lambda:InvokeFunction", + Principal='432143214321', + SourceArn="arn:aws:lambda:us-west-2:account-id:function:helloworld", + SourceAccount='123412341234', + EventSourceToken='blah', + Qualifier='2' + ) + assert 'Statement' in response + res = json.loads(response['Statement']) + assert res['Action'] == "lambda:InvokeFunction" + + @mock_lambda def get_function_policy(): conn = boto3.client('lambda', 'us-west-2') @@ -664,6 +697,17 @@ def get_function_policy(): Publish=True, ) + response = conn.add_permission( + FunctionName='testFunction', + StatementId='1', + Action="lambda:InvokeFunction", + Principal='432143214321', + SourceArn="arn:aws:lambda:us-west-2:account-id:function:helloworld", + SourceAccount='123412341234', + EventSourceToken='blah', + Qualifier='2' + ) + response = conn.get_policy( FunctionName='testFunction' ) @@ -671,4 +715,4 @@ def get_function_policy(): assert 'Policy' in response assert isinstance(response['Policy'], str) res = json.loads(response['Policy']) - assert res['Statement'][0]['Action'] == ['lambda:InvokeFunction'] \ No newline at end of file + assert res['Statement'][0]['Action'] == 'lambda:InvokeFunction'