Fix: Handling Invalid events and reset state of local variables

Fix: Handling Invalid events and reset state of local variables

Author: Keyur

moesif_aws_lambda/middleware.py

@@ -52,14 +52,23 @@ def __init__(self, handler):
             self.DEBUG = self.moesif_options.get('DEBUG', False)
             self.event = None
             self.context = None
-            self.start_time = datetime.utcnow()
 
             # Intialized the client
             if os.environ.get("MOESIF_APPLICATION_ID"):
                 self.api_client = MoesifAPIClient(os.environ["MOESIF_APPLICATION_ID"]).api
             else:
                 raise Exception('Moesif Application ID is required in settings')
 
+        def clear_state(self):
+            """Function to clear state of local variable"""
+            self.event = None
+            self.context = None
+            self.event_req = None
+            self.metadata = None
+            self.session_token = None
+            self.user_id = None
+            self.company_id = None
+
         def get_user_id(self, event, context):
             """Function to fetch UserId"""
             username = None
@@ -130,10 +139,19 @@ def process_body(self, body_wrapper):
         def before(self, event, context):
             """This function runs before the handler is invoked, is passed the event & context and must return an event & context too."""
 
+            # Clear the state of the local variables
+            self.clear_state()
+
+            # Set/Save event and context for use Skip Event function
+            self.event = event
+            self.context = context
+
             # Request Method
             request_verb = event.get('httpMethod')
             if request_verb is None:
                 print('MOESIF: [before] AWS Lambda trigger must be a Load Balancer or API Gateway See https://docs.aws.amazon.com/lambda/latest/dg/services-alb.html or https://docs.aws.amazon.com/lambda/latest/dg/with-on-demand-https.html.')
+                self.event = None
+                self.context = None
                 return event, context
 
             # Request headers
@@ -149,9 +167,9 @@ def before(self, event, context):
             # Request Time
             epoch = event and event.get('request_context', {}).get('requestTimeEpoch')
             if epoch is not None: 
-                request_time =datetime.utcfromtimestamp(epoch)
+                request_time = datetime.utcfromtimestamp(epoch)
             else:
-                request_time = self.start_time
+                request_time = datetime.utcnow()
 
             # Request Body
             req_body, req_transfer_encoding = self.process_body(event)
@@ -234,66 +252,64 @@ def before(self, event, context):
                 body = req_body,
                 transfer_encoding = req_transfer_encoding)
 
-            # Set/Save event and context for use Skip Event function
-            self.event = event
-            self.context = context
-
             # Return event, context
             return event, context
 
         def after(self, retval):
             """This function runs after the handler is invoked, is passed the response and must return an response too."""
-            # Response body
-            resp_body, resp_transfer_encoding = self.process_body(retval)
+            
+            if self.event is not None:
+                # Response body
+                resp_body, resp_transfer_encoding = self.process_body(retval)
 
-            # Event Response object
-            event_rsp = EventResponseModel(time = datetime.utcnow().strftime("%Y-%m-%dT%H:%M:%S.%f")[:-3],
-                status = retval.get('statusCode', 599),
-                headers = retval.get('headers', {}),
-                body = resp_body,
-                transfer_encoding = resp_transfer_encoding)
+                # Event Response object
+                event_rsp = EventResponseModel(time = datetime.utcnow().strftime("%Y-%m-%dT%H:%M:%S.%f")[:-3],
+                    status = retval.get('statusCode', 599),
+                    headers = retval.get('headers', {}),
+                    body = resp_body,
+                    transfer_encoding = resp_transfer_encoding)
 
-            # Event object
-            event_model = EventModel(request = self.event_req,
-                response = event_rsp,
-                user_id = self.user_id,
-                company_id = self.company_id,
-                session_token = self.session_token,
-                metadata = self.metadata)
-            
-            # Mask Event Model
-            try:
-                mask_event_model = self.moesif_options.get('MASK_EVENT_MODEL', None)
-                if mask_event_model is not None:
-                    event_model = mask_event_model(event_model)
-            except:
-                if self.DEBUG:
-                    print("MOESIF Can not execute MASK_EVENT_MODEL function. Please check moesif settings.")
+                # Event object
+                event_model = EventModel(request = self.event_req,
+                    response = event_rsp,
+                    user_id = self.user_id,
+                    company_id = self.company_id,
+                    session_token = self.session_token,
+                    metadata = self.metadata)
+                
+                # Mask Event Model
+                try:
+                    mask_event_model = self.moesif_options.get('MASK_EVENT_MODEL', None)
+                    if mask_event_model is not None:
+                        event_model = mask_event_model(event_model)
+                except:
+                    if self.DEBUG:
+                        print("MOESIF Can not execute MASK_EVENT_MODEL function. Please check moesif settings.")
 
-            # Skip Event
-            try:
-                skip_event = self.moesif_options.get('SKIP', None)
-                if skip_event is not None:
-                    if skip_event(self.event, self.context):
-                        if self.DEBUG:
-                            print('MOESIF Skip sending event to Moesif')
-                        return retval
-            except:
-                if self.DEBUG:
-                    print("MOESIF Having difficulty executing skip_event function. Please check moesif settings.")
+                # Skip Event
+                try:
+                    skip_event = self.moesif_options.get('SKIP', None)
+                    if skip_event is not None:
+                        if skip_event(self.event, self.context):
+                            if self.DEBUG:
+                                print('MOESIF Skip sending event to Moesif')
+                            return retval
+                except:
+                    if self.DEBUG:
+                        print("MOESIF Having difficulty executing skip_event function. Please check moesif settings.")
 
-            # Add direction field
-            event_model.direction = "Incoming"
+                # Add direction field
+                event_model.direction = "Incoming"
+
+                # Send event to Moesif
+                if self.DEBUG:
+                    print('Moesif Event Model:')
+                    print(json.dumps(self.event))
+                
+                event_send = self.api_client.create_event(event_model)
+                if self.DEBUG:
+                    print('MOESIF ' + str(event_send))
 
-            # Send event to Moesif
-            if self.DEBUG:
-                print('Moesif Event Model:')
-                print(json.dumps(self.event))
-            
-            event_send = self.api_client.create_event(event_model)
-            if self.DEBUG:
-                print('MOESIF ' + str(event_send))
-            
             # Send response
             return retval