adding support for POST method and applying limit in payload lenght

Author: GGontijo

elasticsearch_middleware.py

@@ -25,37 +25,30 @@ def __init__(self, app: ASGIApp, config: dict) -> None:
             config (dict): Configuration settings for Elasticsearch logging.
                 {
                     'url': str,          # Elasticsearch server URL
+                    'user': str,         # Elasticsearch API user
+                    'password': str      # Elasticsearch API password
                     'index': str,        # Elasticsearch index name
                     'environment': str,  # Environment identifier for logs
+                    'limit': bool,       # Limit Elasticsearch payload array and string lenght
                     'debug': bool        # When True logs aren't sent to Elasticsearch
                 }
         """
         elastic_config = config
-        self.elasticsearch_client = Elasticsearch([elastic_config.get('url')])
+        elastic_user = config.get('user')
+        elastic_password = config.get('password')
+        basic_auth = (elastic_user, elastic_password) if elastic_user and elastic_password else None
+        self.elasticsearch_client = Elasticsearch([elastic_config.get('url')], basic_auth=basic_auth)
         self.index = elastic_config.get('index')
-        self.environment = elastic_config.get('environment')
-        self.debug = elastic_config.get('debug')
+        self.limit = elastic_config.get('limit', False)
+        self.environment = elastic_config.get('environment', 'Development')
+        self.debug = elastic_config.get('debug', False)
         self.app = app
 
-    async def __call__(self, scope: Scope, receive: Receive, send: Send):
+    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
+
         if scope.get('type') == 'http' and not self.debug:
-            request = Request(scope)
-            
-            request_data_received = await receive() 
-            request_received_body = json.dumps(json.loads(request_data_received.get('body')), ensure_ascii=False) if len(request_data_received.get('body')) > 0 else None
-            
-            start_time = time.time()
-            
-            log_data = {
-                "@timestamp": datetime.utcnow().replace(tzinfo=timezone.utc).strftime("%Y-%m-%dT%H:%M:%S"),
-                "environment": self.environment, 
-                "method": request.method,
-                "path": request.url.path,
-                "request": {},
-                "response": {}
-            }
 
-            log_data["request"]["body"] = request_received_body
+            request = Request(scope)
 
             async def intercept_send(response):
                 nonlocal log_data
@@ -64,22 +57,26 @@ async def intercept_send(response):
                     return 
                 if response['type'] == 'http.response.body' and "response" in log_data.keys(): # Response part
 
+                    # Finishes telemetry
                     end_time = time.time()
                     elapsed_time = end_time - start_time
 
                     log_data["elapsed_time"] = elapsed_time
 
                     if log_data["response"]["headers"].get('content-type') == 'application/json':
-                        log_data["response"]["body"] = json.dumps(json.loads(response.get('body')), ensure_ascii=False) if 'body' in response.keys() else None
-                        
+                        response_body = json.loads(response.get('body'))
+                        response_body = self.limit_array_length(response_body)
+
+                        response_body = json.dumps(response_body, ensure_ascii=False) if 'body' in response.keys() else None
+                        log_data["response"]["body"] = response_body
+
                     elif log_data["response"]["headers"].get('content-type') == 'application/octet-stream':
                         log_data["response"]["body"] = str(response.get('body')) if 'body' in response.keys() else None
 
                     self.log_to_elasticsearch(log_data)
 
                 if response['type'] == 'http.response.start': # Request part
 
-
                     request_headers = dict(request.headers) if 'headers' in request.keys() else None
                     request_query_parameters = dict(request.query_params) if len(request.query_params._list) > 0 else None
 
@@ -93,13 +90,80 @@ async def intercept_send(response):
 
                 await send(response)
 
-            await self.app(scope, receive, intercept_send)
+            start_time = time.time()
+            
+            log_data = {
+                "@timestamp": datetime.utcnow().replace(tzinfo=timezone.utc).strftime("%Y-%m-%dT%H:%M:%S"),
+                "environment": self.environment, 
+                "method": request.method,
+                "path": request.url.path,
+                "request": {},
+                "response": {}
+            }
+ 
+            # Starts telemetry
+            start_time = time.time()
+
+            async def intercept_receive():
+                message = await receive()
+
+                more_body = message.get("more_body", False)
+                body = message.get("body", "")
+                while more_body:
+                    message = await receive()
+                    body += message.get("body", b"")
+                    more_body = message.get("more_body", False)
+
+                message["body"] = body
+                request_body = ''
+
+                if len(message["body"]) > 0:
+                    request_body = json.loads(body.decode('utf-8'))
+                    request_body = self.limit_string_length(request_body)
+                    request_body = json.dumps(request_body, ensure_ascii=False) if len(body.decode('utf-8')) > 0 else None
+                
+                log_data["request"]["body"] = request_body
+
+                return message
+        
+            await self.app(scope, intercept_receive, intercept_send)
+            
         else:
             await self.app(scope, receive, send)
 
-    def log_to_elasticsearch(self, log_data):
+    def log_to_elasticsearch(self, log_data) -> None:
         try:
             self.elasticsearch_client.index(index=self.index, body=log_data)
             log_data.clear()
         except Exception as e:
             logging.error(f"Failed to log to Elasticsearch: {str(e)}")
+
+    def limit_string_length(self, data, max_lines=50):
+        if not self.limit:
+            return data
+        if isinstance(data, dict):
+            for key, value in data.items():
+                data[key] = self.limit_string_length(value, max_lines)
+        elif isinstance(data, list):
+            for i in range(len(data)):
+                data[i] = self.limit_string_length(data[i], max_lines)          
+        elif isinstance(data, str) and len(data.split('\n')) > max_lines:
+            data_splitted = data.split('\n')[:max_lines]
+            data_splitted.append(f' [...] value limited in {max_lines} lines')
+            data = '\n'.join(data_splitted)
+        elif isinstance(data, str) and len(data.split('/')) > max_lines: # Base64 files
+            data_splitted = data.split('/')[:max_lines]
+            data_splitted.append(f' [...] value limited in {max_lines} lines')
+            data = '/'.join(data_splitted)
+        return data
+    
+    def limit_array_length(self, data, max_length=3):
+        if not self.limit:
+            return data
+        if isinstance(data, dict):
+            for key, value in data.items():
+                data[key] = self.limit_array_length(value, max_length)
+        elif isinstance(data, list) and len(data) > max_length:
+            data = data[:max_length]
+            data.append(f'[...] array limited in {max_length} objects')
+        return data