Merge pull request #38 from cloudblue/LITE-25143_improve_logging

LITE-25143 obfuscate sensitive data

connect/eaas/core/proto.py

@@ -1,6 +1,32 @@
 from typing import Any, List, Literal, Optional, Union
 
-from pydantic import BaseModel, Field
+from pydantic import BaseModel as PydanticBaseModel, Field
+from pydantic.utils import DUNDER_ATTRIBUTES
+
+
+class BaseModel(PydanticBaseModel):
+
+    def get_sensitive_fields(self):
+        return []
+
+    def __obfuscate_args__(self, k, v):
+        if isinstance(v, str) and v and k in self.get_sensitive_fields():
+            return k, f'{v[0:2]}******{v[-2:]}'
+
+        if isinstance(v, (list, dict)) and v:
+            return k, '******'
+
+        return k, v
+
+    def __repr_args__(self):
+        return [
+            self.__obfuscate_args__(k, v)
+            for k, v in self.__dict__.items()
+            if (
+                k not in DUNDER_ATTRIBUTES
+                and (k not in self.__fields__ or self.__fields__[k].field_info.repr)
+            )
+        ]
 
 
 class MessageType:
@@ -25,6 +51,9 @@ class TaskOptions(BaseModel):
     installation_id: Optional[str]
     connect_correlation_id: Optional[str]
 
+    def get_sensitive_fields(self):
+        return ['api_key']
+
 
 class TaskOutput(BaseModel):
     result: str
@@ -62,6 +91,9 @@ class Logging(BaseModel):
     runner_log_level: Optional[str]
     meta: Optional[LogMeta]
 
+    def get_sensitive_fields(self):
+        return ['logging_api_key']
+
 
 class EventDefinition(BaseModel):
     event_type: str
@@ -102,6 +134,9 @@ class SetupRequest(BaseModel):
     runner_version: Optional[str]
     model_type: Literal['setup_request'] = 'setup_request'
 
+    def get_sensitive_fields(self):
+        return ['variables']
+
 
 class HttpResponse(BaseModel):
     status: int
@@ -115,6 +150,9 @@ class HttpRequest(BaseModel):
     headers: dict
     content: Optional[Any]
 
+    def get_sensitive_fields(self):
+        return ['headers']
+
 
 class WebTaskOptions(BaseModel):
     correlation_id: str
@@ -128,6 +166,9 @@ class WebTaskOptions(BaseModel):
     call_type: Optional[Literal['admin', 'user']]
     call_source: Optional[Literal['ui', 'api']]
 
+    def get_sensitive_fields(self):
+        return ['api_key']
+
 
 class WebTask(BaseModel):
     options: WebTaskOptions

tests/connect/eaas/core/test_proto.py

@@ -4,6 +4,7 @@
 
 from connect.eaas.core.proto import (
     EventDefinition,
+    HttpRequest,
     Logging,
     LogMeta,
     Message,
@@ -15,6 +16,7 @@
     Task,
     TaskInput,
     TaskOptions,
+    WebTaskOptions,
 )
 
 TASK_DATA = {
@@ -442,3 +444,88 @@ def test_serialize_setup_response_with_events_definitions():
     assert serialized['version'] == 2
     assert serialized['message_type'] == MessageType.SETUP_RESPONSE
     assert serialized['data'] == data
+
+
+def test_obfuscate_task_options():
+    task_options = TaskOptions(
+        task_id='task_id',
+        task_category='task_category',
+        api_key='This is my API key',
+    )
+
+    fields = {}
+    representation = str(task_options).split(' ')
+    for field in representation:
+        k, v = field.split('=')
+        fields[k] = v[1:-1] if isinstance(v, str) else v
+
+    assert fields['api_key'] == 'Th******ey'
+
+
+def test_obfuscate_logging():
+    logging = Logging(
+        logging_api_key='This is my API key',
+    )
+
+    fields = {}
+    representation = str(logging).split(' ')
+    for field in representation:
+        k, v = field.split('=')
+        fields[k] = v[1:-1] if isinstance(v, str) else v
+
+    assert fields['logging_api_key'] == 'Th******ey'
+
+
+def test_obfuscate_setup_response():
+    setup_response = SetupResponse(variables={'VAR1': 'VAL1'})
+
+    fields = {}
+    representation = str(setup_response).split(' ')
+    for field in representation:
+        k, v = field.split('=')
+        fields[k] = v[1:-1] if isinstance(v, str) else v
+
+    assert fields['variables'] == '******'
+
+
+def test_obfuscate_setup_request():
+    setup_request = SetupRequest(variables=[{'VAR1': 'VAL1'}])
+
+    fields = {}
+    representation = str(setup_request).split(' ')
+    for field in representation:
+        k, v = field.split('=')
+        fields[k] = v[1:-1] if isinstance(v, str) else v
+
+    assert fields['variables'] == '******'
+
+
+def test_obfuscate_http_request():
+    http_request = HttpRequest(
+        method='method',
+        url='url',
+        headers={'Authorization': 'My Api Key'},
+    )
+
+    fields = {}
+    representation = str(http_request).split(' ')
+    for field in representation:
+        k, v = field.split('=')
+        fields[k] = v[1:-1] if isinstance(v, str) else v
+
+    assert fields['headers'] == '******'
+
+
+def test_obfuscate_webtask_options():
+    webtask_options = WebTaskOptions(
+        correlation_id='correlation_id',
+        reply_to='reply_to',
+        api_key='This is my API key',
+    )
+    fields = {}
+    representation = str(webtask_options).split(' ')
+    for field in representation:
+        k, v = field.split('=')
+        fields[k] = v[1:-1] if isinstance(v, str) else v
+
+    assert fields['api_key'] == 'Th******ey'