Feature/data redaction (#6)

Author: gokhan721

Pipfile

@@ -7,6 +7,7 @@ name = "pypi"
 six = "*"
 websocket-client = "*"
 pystache = "*"
+cachetools= "*"
 
 [dev-packages]
 

setup.py

@@ -73,7 +73,8 @@ def ReadConfig(section, value, default):
 
 install_requires = ['six >= 1.11',
                     'websocket-client >= 0.56.0',
-                    'pystache >= 0.6.0'
+                    'pystache >= 0.6.0',
+                    'cachetools >= 5.2.0'
                     ]
 
 if sys.version_info[0] == 2:

tracepointdebug/__init__.py

@@ -4,21 +4,28 @@
 from .broker.broker_manager import BrokerManager
 from .probe.trace_point_manager import TracePointManager
 from .probe.log_point_manager import LogPointManager
+from .probe.error_stack_manager import ErrorStackManager
+
 '''
     After importing ConfigProvider for the first time, the __init__.py has been run by interpreter and
     whole configuration is reflected to configs.
 '''
-from .config.config_provider import ConfigProvider
 
 
+tracepoint_data_redaction_callback = None
+log_data_redaction_callback = None
+
 import logging
 logger = logging.getLogger(__name__)
 
-def start():
+def start(tracepoint_data_redaction_callback=None, log_data_redaction_callback=None):
     cdbg_native.InitializeModule(None)
     _broker_manager = BrokerManager().instance()
+    _broker_manager.initialize(tracepoint_data_redaction_callback, log_data_redaction_callback)
     tpm = TracePointManager(broker_manager=_broker_manager)
     lpm = LogPointManager(broker_manager=_broker_manager)
-    _broker_manager.initialize()
+    esm = ErrorStackManager(broker_manager=_broker_manager)
+    esm.start()
     atexit.register(tpm.remove_all_trace_points)
     atexit.register(lpm.remove_all_log_points)
+    atexit.register(esm.shutdown)

tracepointdebug/broker/broker_client.py

@@ -5,6 +5,8 @@
 from time import sleep
 
 import websocket
+from tracepointdebug.config import config_names
+from tracepointdebug.config.config_provider import ConfigProvider
 
 from tracepointdebug.utils import debug_logger
 from tracepointdebug.broker.ws_app import WSApp
@@ -38,6 +40,7 @@ def __init__(self, host, port, broker_credentials, message_callback, initial_req
         self.connection_timer = None
         self.connection_timeout = 10
         self.reconnect_interval = 3
+        self.error_printed = False
         self.connected = threading.Event()
         self.initial_request_to_broker = initial_request_to_broker
 
@@ -133,25 +136,37 @@ def on_error(self, ws, msg):
                 self._running = False
                 if self.ws:
                     self.ws.close()
-        logger.error("Error on connection, msg: {}".format(msg))
+        if not self.error_printed:
+            logger.error("Error on connection, msg: {}".format(msg))
+            self.error_printed = True
+        else:
+            debug_logger("Error on connection, msg: {}".format(msg))
 
     def on_close(self, ws):
+        self.error_printed = False
         debug_logger("Connection closed")
 
     def on_open(self, ws):
         debug_logger("Connection open")
+        self.error_printed = False
         self.connected.set()
         connection_set = self.connected.wait() #TODO Timeout
         if connection_set:
             self.initial_request_to_broker()
 
     def send(self, data):
         try:
-            self.ws.send(data)
+            if self.ws.sock.connected:
+                self.ws.send(data)
+            else:
+                if ConfigProvider.get(config_names.SIDEKICK_PRINT_CLOSED_SOCKET_DATA, False):
+                    print("Socket is already closed while sending data: %s" % data)
+                debug_logger("Socket is already closed while sending data to see data set SIDEKICK_PRINT_DEBUG_DATA to True!")
         except websocket.WebSocketConnectionClosedException as e:
-            logger.error("Error sending %s" % e)
+            debug_logger("Error sending %s" % e)
 
     def close(self):
+        self.error_printed = False
         self._running = False
         if self.ws:
             self.ws.close()

tracepointdebug/broker/broker_manager.py

@@ -43,6 +43,8 @@ def __init__(self):
         self.initialized = False
         self._event_executor = ThreadPoolExecutor()
         self._request_executor = ThreadPoolExecutor()
+        self._tracepoint_data_redaction_callback = None
+        self._log_data_redaction_callback = None
         import sys
         if sys.version_info[0] >= 3:
             self.application_status_thread = Thread(target=self.application_status_sender, daemon=True)
@@ -57,10 +59,14 @@ def instance():
         return BrokerManager() if BrokerManager.__instance is None else BrokerManager.__instance
 
 
-    def initialize(self):
+    def initialize(self, tracepoint_data_redaction_callback=None, log_data_redaction_callback=None):
         if not self.initialized:
             self.connect_to_broker()
             self.initialized = True
+            if callable(tracepoint_data_redaction_callback):
+                self._tracepoint_data_redaction_callback = tracepoint_data_redaction_callback
+            if callable(log_data_redaction_callback):
+                self._log_data_redaction_callback = log_data_redaction_callback
 
 
     def connect_to_broker(self):
@@ -98,7 +104,6 @@ def prepare_event(event):
         event.application_instance_id = application_info['applicationInstanceId']
         event.application_name = application_info['applicationName']
 
-
     def do_publish_event(self, event):
         self.prepare_event(event)
         try:
@@ -168,6 +173,5 @@ def publish_application_status(self, client=None):
 
         for status_provider in self.application_status_providers:
             status_provider.provide(application_status, client)
-
         event = ApplicationStatusEvent(client=client, application=application_status)
         self.publish_event(event)

tracepointdebug/broker/broker_message_callback.py

@@ -10,6 +10,7 @@
 from tracepointdebug.probe.handler.request.logPoint.remove_log_point_request_handler import RemoveLogPointRequestHandler
 from tracepointdebug.probe.handler.request.logPoint.update_log_point_request_handler import UpdateLogPointRequestHandler
 from tracepointdebug.probe.handler.response.filter_logpoints_response_handler import FilterLogPointsResponseHandler
+from tracepointdebug.utils import debug_logger
 
 MESSAGE_REQUEST_TYPE = "Request"
 MESSAGE_RESPONSE_TYPE = "Response"
@@ -50,16 +51,15 @@ def on_message(self, broker_client, message):
                     serialized = to_json(response)
                     broker_client.send(serialized)
                 else:
-                    print("No request handler could be found for message with name {}: {}".format(message.get("name"),
+                    debug_logger("No request handler could be found for message with name {}: {}".format(message.get("name"),
                                                                                               message))
             elif message_type == MESSAGE_RESPONSE_TYPE:
                 handler = RESPONSE_HANDLER_MAP.get(message.get("name"))
                 if handler is not None:
                     response = handler.get_response_cls()(**message)
                     handler.handle_response(response)
                 else:
-                    print("No response handler could be found for message with name {}: {}".format(message.get("name"),
+                    debug_logger("No response handler could be found for message with name {}: {}".format(message.get("name"),
                                                                                               message))
-
         except Exception as e:
-            print(e)
+            debug_logger(e)

tracepointdebug/config/config_metadata.py

@@ -8,6 +8,18 @@
         'type': 'boolean',
         'defaultValue': False,
     },
+    config_names.SIDEKICK_ERROR_STACK_ENABLE: {
+        'type': 'boolean',
+        'defaultValue': False,
+    },
+    config_names.SIDEKICK_ERROR_FRAME_COLLECTION_ENABLE: {
+        'type': 'boolean',
+        'defaultValue': False,
+    },
+    config_names.SIDEKICK_PRINT_CLOSED_SOCKET_DATA: {
+        'type': 'boolean',
+        'defaultValue': False,
+    },
     config_names.SIDEKICK_APPLICATION_ID: {
         'type': 'string',
     },

tracepointdebug/config/config_names.py

@@ -12,4 +12,7 @@
 SIDEKICK_APPLICATION_CLASS_NAME = 'sidekick.application.classname'
 SIDEKICK_APPLICATION_VERSION = 'sidekick.application.version'
 SIDEKICK_APPLICATION_TAG_PREFIX = 'sidekick.application.tag.prefix'
-SIDEKICK_APPLICATION_REGION = 'sidekick.application.region'
+SIDEKICK_APPLICATION_REGION = 'sidekick.application.region'
+SIDEKICK_ERROR_STACK_ENABLE = 'sidekick.error.stack.enable'
+SIDEKICK_ERROR_FRAME_COLLECTION_ENABLE = 'sidekick.error.frame.collection.enable'
+SIDEKICK_PRINT_CLOSED_SOCKET_DATA = 'sidekick.print.closed.socket.data'

tracepointdebug/probe/encoder.py

@@ -1,5 +1,5 @@
 import json
-
+from tracepointdebug.utils import debug_logger
 
 class JSONEncoder(json.JSONEncoder):
     def default(self, z):
@@ -11,7 +11,7 @@ def default(self, z):
             else:
                 return super(JSONEncoder, self).default(z)
         except Exception as e:
-            print(e)
+            debug_logger(e)
 
 
 def to_json(data, separators=None):

tracepointdebug/probe/error_stack_manager.py

@@ -0,0 +1,123 @@
+import time, traceback
+from tracepointdebug.probe.coded_exception import CodedException
+from tracepointdebug.probe.coded_exception import CodedException
+from tracepointdebug.probe.event.errorstack.error_stack_rate_limit_event import ErrorStackRateLimitEvent
+from tracepointdebug.probe.event.errorstack.error_stack_snapshot_event import ErrorStackSnapshotEvent
+from tracepointdebug.probe.event.errorstack.error_stack_snapshot_failed_event import ErrorStackSnapshotFailedEvent
+from tracepointdebug.probe.ratelimit.rate_limit_result import RateLimitResult
+from tracepointdebug.probe.ratelimit.rate_limiter import RateLimiter
+from tracepointdebug.probe.snapshot_collector import SnapshotCollector
+import logging, sys, threading
+from tracepointdebug.config import config_names
+from tracepointdebug.config.config_provider import ConfigProvider
+from datetime import datetime as dt
+from cachetools import TTLCache
+import datetime, os
+
+logger = logging.getLogger(__name__)
+
+_MAX_SNAPSHOT_SIZE = 32768
+_MAX_FRAMES = 10
+_MAX_EXPAND_FRAMES = 2
+_MAX_TIME_TO_ALIVE_MIN = 5
+
+class ErrorStackManager(object):
+    __instance = None
+
+    def __init__(self, broker_manager):
+        self.broker_manager = broker_manager
+        self.old_settrace = sys.gettrace()
+        self.old_threading = threading._trace_hook
+        self.condition = None
+        self.timer = None
+        self.sidekick_exception = "sidekickException"
+        self.rate_limiter = RateLimiter()
+        self.ttl_cache = TTLCache(maxsize=2048, ttl=datetime.timedelta(minutes=_MAX_TIME_TO_ALIVE_MIN), timer=datetime.datetime.now)
+        ErrorStackManager.__instance = self
+
+    @staticmethod
+    def instance(*args, **kwargs):
+        return ErrorStackManager(*args,
+                                 **kwargs) if ErrorStackManager.__instance is None else ErrorStackManager.__instance
+
+    @staticmethod
+    def get_id(file, line):
+        return '{}:{}:{}'.format(file, line, str(dt.now()))
+
+    def _get_point_cache_id(self, frame):
+        return frame.f_code.co_filename + ":::" + str(frame.f_lineno)
+
+    def _check_point_inserted(self, frame):
+        error_point_id = self._get_point_cache_id(frame)
+        item = self.ttl_cache.get(error_point_id, None)
+        if item is None:
+            self.ttl_cache[error_point_id] = True
+            return False
+        return True
+
+    def _white_list_exceptions(self, frame):
+        frame_file_path = os.path.abspath(frame.f_code.co_filename)
+        blacklist = ["python", "site-packages", "importlib", "tracepointdebug"]
+        for black in blacklist:
+            if black in frame_file_path:
+                return False
+        return True
+
+    def trace_hook(self, frame, event, arg):
+        if not self._white_list_exceptions(frame):
+            return
+        frame.f_trace = self._frame_hook
+
+    def _frame_hook(self, frame, event, arg):
+        try:
+            if event != "exception":
+                return
+            frame_file_name = frame.f_code.co_filename
+            frame_line_no = frame.f_lineno
+            rate_limit_result_for_frame_call = self.rate_limiter.check_rate_limit(time.time())
+            check_point_already_inserted = self._check_point_inserted(frame)
+
+            if (check_point_already_inserted):
+                return
+
+            if (rate_limit_result_for_frame_call == RateLimitResult.HIT):
+                event = ErrorStackRateLimitEvent(frame_file_name, frame_line_no)
+                self._publish_event(event)
+
+            if (rate_limit_result_for_frame_call == RateLimitResult.EXCEEDED):
+                return
+
+            frames = []
+            if ConfigProvider.get(config_names.SIDEKICK_ERROR_FRAME_COLLECTION_ENABLE, False):
+                snapshot_collector = SnapshotCollector(_MAX_SNAPSHOT_SIZE, _MAX_FRAMES, _MAX_EXPAND_FRAMES)
+                snapshot = snapshot_collector.collect(frame)
+                frames = snapshot.frames
+            error_stack_id = self.get_id(frame_file_name, frame_line_no)
+            error = {
+                "name": str(arg[0]) or "Error",
+                "message": str(arg[1]),
+                "stack": str(traceback.extract_tb(arg[2]))
+            }
+            event = ErrorStackSnapshotEvent(error_stack_id, frame_file_name, frame_line_no, method_name=frame.f_code.co_name,
+                                            error=error, frames=frames)
+            self._publish_event(event)
+        except Exception as exc:
+            logger.warning('Error on error stack snapshot %s' % exc)
+            code = 0
+            if isinstance(exc, CodedException):
+                code = exc.code
+            event = ErrorStackSnapshotFailedEvent(frame.f_code.co_filename, frame.f_lineno, code, str(exc))
+            self._publish_event(event)
+
+    def start(self):
+        if ConfigProvider.get(config_names.SIDEKICK_ERROR_STACK_ENABLE):
+            sys.settrace(self.trace_hook)
+            threading.settrace(self.trace_hook)
+
+    def shutdown(self):
+        if ConfigProvider.get(config_names.SIDEKICK_ERROR_STACK_ENABLE):
+            sys.settrace(self.old_settrace)
+            threading.settrace(self.old_threading)
+
+    def _publish_event(self, event):
+        self.broker_manager.publish_event(event)