From 0f9ce4ab5df9c75497d63bc2dbcec89a3bea3cb2 Mon Sep 17 00:00:00 2001
From: Torgayev Tamirlan <10016756+prog893@users.noreply.github.com>
Date: Tue, 7 Jan 2020 09:26:41 +0600
Subject: [PATCH] Add top-level JSON type checks (#6)

* Add top-level JSON type checks

* Fix logging
---
 lambda/main.py | 28 ++++++++++++++++++----------
 main.tf        |  2 +-
 2 files changed, 19 insertions(+), 11 deletions(-)

diff --git a/lambda/main.py b/lambda/main.py
index 7f84667..210cdc8 100644
--- a/lambda/main.py
+++ b/lambda/main.py
@@ -84,29 +84,31 @@ def append_to_dict(dictionary: dict, log_type: str, log_data: object, log_timest
     dictionary[log_type]['records'].append(log_data)
 
 
-def normalize_kinesis_payload(payload: dict):
+def normalize_kinesis_payload(p: dict):
     # Normalize messages from CloudWatch (subscription filters) and pass through anything else
     # https://docs.aws.amazon.com/ja_jp/AmazonCloudWatch/latest/logs/SubscriptionFilters.html
 
-    logger.debug(f"normalizer input: {payload}")
+    logger.debug(f"normalizer input: {p}")
 
     payloads = []
 
-    if len(payload) < 1:
-        logger.error(f"Got weird record: \"{payload}\", skipping")
+    if len(p) < 1:
+        logger.error(f"Got weird record: \"{p}\", skipping")
         return payloads
 
     # check if data is JSON and parse
     try:
-        payload = json.loads(payload)
+        payload = json.loads(p)
+        if type(payload) is not dict:
+            logger.error(f"Top-level JSON data is not an object, giving up: {payload}")
+            return payloads
 
     except JSONDecodeError:
-        logger.error(f"Non-JSON data found: {payload}, skipping")
+        logger.error(f"Non-JSON data found: {p}, giving up")
         return payloads
 
     if 'messageType' in payload:
-        logger.debug(f"Got payload looking like CloudWatch Logs via subscription filters: "
-                     f"{payload}")
+        logger.debug(f"Got payload looking like CloudWatch Logs via subscription filters: {payload}")
 
         if payload['messageType'] == "DATA_MESSAGE":
             if 'logEvents' in payload:
@@ -117,7 +119,13 @@ def normalize_kinesis_payload(payload: dict):
                         payload_parsed = json.loads(event['message'])
                         logger.debug(f"parsed payload: {payload_parsed}")
 
-                    except JSONDecodeError:
+                        if type(payload_parsed) is not dict:
+                            logger.error(f"Top-level JSON data in CWL payload is not an object, giving up: "
+                                         f"{payload_parsed}")
+                            continue
+
+                    except JSONDecodeError as e:
+                        logger.debug(e)
                         logger.debug(f"Non-JSON data found inside CWL message: {event}, giving up")
                         continue
 
@@ -136,7 +144,7 @@ def normalize_kinesis_payload(payload: dict):
             raise ValueError(f"Unknown messageType: {payload}")
     else:
         payloads.append(payload)
-        
+
     return payloads
 
 
diff --git a/main.tf b/main.tf
index 541651f..013d29c 100644
--- a/main.tf
+++ b/main.tf
@@ -77,7 +77,7 @@ resource "aws_iam_role_policy_attachment" "xray_access" {
 
 module "iam" {
   source  = "baikonur-oss/iam-nofile/aws"
-  version = "1.0.1"
+  version = "v2.0.0"
 
   type = "lambda"
   name = var.name