From beff8a0dacdf64a69dca7acd5f85338ec14ddae8 Mon Sep 17 00:00:00 2001
From: Jente Peeraer <jente.peeraer@cloudway.be>
Date: Mon, 8 Jul 2019 10:21:06 +0200
Subject: [PATCH] avoid keyerrors by using the get function on an object

---
 .gitignore           |  2 ++
 flask_lambda.py      | 19 ++++++++++---------
 util/__init__.py     |  1 +
 util/object_utils.py |  9 +++++++++
 4 files changed, 22 insertions(+), 9 deletions(-)
 create mode 100644 util/__init__.py
 create mode 100644 util/object_utils.py

diff --git a/.gitignore b/.gitignore
index 72364f9..9719cdf 100644
--- a/.gitignore
+++ b/.gitignore
@@ -87,3 +87,5 @@ ENV/
 
 # Rope project settings
 .ropeproject
+
+.idea
\ No newline at end of file
diff --git a/flask_lambda.py b/flask_lambda.py
index 41db471..2d59340 100644
--- a/flask_lambda.py
+++ b/flask_lambda.py
@@ -32,6 +32,7 @@
         from io import StringIO
 
 from flask import Request
+from util import get_nested
 
 
 __version__ = '0.1.2'
@@ -50,12 +51,12 @@ def make_environ(event, context):
         http_hdr_name = 'HTTP_{}'.format(hdr_name)
         environ[http_hdr_name] = hdr_value
 
-    qs = event['queryStringParameters']
+    qs = event.get('queryStringParameters', '')
 
-    environ['REQUEST_METHOD'] = event['httpMethod']
-    environ['PATH_INFO'] = event['path']
+    environ['REQUEST_METHOD'] = event.get('httpMethod', '')
+    environ['PATH_INFO'] = event.get('path', '')
     environ['QUERY_STRING'] = urlencode(qs) if qs else ''
-    environ['REMOTE_ADDR'] = event['requestContext']['identity']['sourceIp']
+    environ['REMOTE_ADDR'] = get_nested(event, '', 'requestContext', 'identity', 'sourceIp')
     environ['HOST'] = '{}:{}'.format(
         environ.get('HTTP_HOST', ''),
         environ.get('HTTP_X_FORWARDED_PORT', ''),
@@ -63,15 +64,15 @@ def make_environ(event, context):
     environ['SCRIPT_NAME'] = ''
     environ['SERVER_NAME'] = 'SERVER_NAME'
 
-    environ['SERVER_PORT'] = environ['HTTP_X_FORWARDED_PORT']
+    environ['SERVER_PORT'] = environ.get('HTTP_X_FORWARDED_PORT', '')
     environ['SERVER_PROTOCOL'] = 'HTTP/1.1'
 
     environ['CONTENT_LENGTH'] = str(
-        len(event['body']) if event['body'] else ''
+        len(event.get('body', ''))
     )
 
-    environ['wsgi.url_scheme'] = environ['HTTP_X_FORWARDED_PROTO']
-    environ['wsgi.input'] = StringIO(event['body'] or '')
+    environ['wsgi.url_scheme'] = environ.get('HTTP_X_FORWARDED_PROTO', '')
+    environ['wsgi.input'] = StringIO(event.get('body', ''))
     environ['wsgi.version'] = (1, 0)
     environ['wsgi.errors'] = sys.stderr
     environ['wsgi.multithread'] = False
@@ -140,5 +141,5 @@ def __call__(self, event, context):
             return {
                 'statusCode': 500,
                 'headers': {},
-                'body': 'internal server error'
+                'body': 'Internal Server Error'
             }
diff --git a/util/__init__.py b/util/__init__.py
new file mode 100644
index 0000000..65242d4
--- /dev/null
+++ b/util/__init__.py
@@ -0,0 +1 @@
+from object_utils import get_nested
\ No newline at end of file
diff --git a/util/object_utils.py b/util/object_utils.py
new file mode 100644
index 0000000..50b6b0e
--- /dev/null
+++ b/util/object_utils.py
@@ -0,0 +1,9 @@
+def get_nested(o, default, *args):
+    if o is None:
+        return default
+    current = o
+    for arg in args:
+        if current is None or arg is None or current.get(arg, None) is None:
+            return default
+        current = current.get(arg, default)
+    return current