From 9bb07e6b6e42811fa768849618a7c32d9f55d4d1 Mon Sep 17 00:00:00 2001
From: Chris Keogh <chris.keogh@xero.com>
Date: Tue, 3 Oct 2017 15:23:00 +1300
Subject: [PATCH] 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 d22d1a7f..935abbcd 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 94c381f5..5215f63c 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 163fa306..317e9f4a 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'