diff --git a/oletools/common/log_helper/_json_formatter.py b/oletools/common/log_helper/_json_formatter.py index 4c5e337..6adc55c 100644 --- a/oletools/common/log_helper/_json_formatter.py +++ b/oletools/common/log_helper/_json_formatter.py @@ -13,8 +13,13 @@ class JsonFormatter(logging.Formatter): Since we don't buffer messages, we always prepend messages with a comma to make the output JSON-compatible. The only exception is when printing the first line, so we need to keep track of it. + + We assume that all input comes from the OletoolsLoggerAdapter which + ensures that there is a `type` field in the record. Otherwise will have + to add a try-except around the access to `record.type`. """ json_dict = dict(msg=record.msg, level=record.levelname) + json_dict['type'] = record.type formatted_message = ' ' + json.dumps(json_dict) if self._is_first_line: diff --git a/oletools/common/log_helper/_logger_adapter.py b/oletools/common/log_helper/_logger_adapter.py index 75e331a..ee52b26 100644 --- a/oletools/common/log_helper/_logger_adapter.py +++ b/oletools/common/log_helper/_logger_adapter.py @@ -8,18 +8,45 @@ class OletoolsLoggerAdapter(logging.LoggerAdapter): """ _json_enabled = None - def print_str(self, message): + def print_str(self, message, **kwargs): """ This function replaces normal print() calls so we can format them as JSON when needed or just print them right away otherwise. """ if self._json_enabled and self._json_enabled(): # Messages from this function should always be printed, - # so when using JSON we log using the same level that set - self.log(_root_logger_wrapper.level(), message) + # so when using JSON we log using the same level that set. + # Additional information in kwargs is added to LogRecord + self.log(_root_logger_wrapper.level(), message, extra=kwargs) else: print(message) + def log(self, lvl, msg, *args, **kwargs): + """ + Run :py:meth:`process` on kwargs, then forward to actual logger. + + This is based on the logging cookbox, section "Using LoggerAdapter to + impart contextual information". + """ + msg, kwargs = self.process(msg, kwargs) + self.logger.log(lvl, msg, *args, **kwargs) + + def process(self, msg, kwargs): + """ + Ensure `kwargs['extra']['type']` exists, init with given arg `type`. + + The `type` field will be added to the :py:class:`logging.LogRecord` and + is used by the :py:class:`JsonFormatter`. + """ + if 'extra' not in kwargs: + kwargs['extra'] = {} + if 'type' in kwargs: + kwargs['extra']['type'] = kwargs['type'] + del kwargs['type'] # downstream loggers cannot deal with this + if 'type' not in kwargs['extra']: + kwargs['extra']['type'] = 'msg' # type will be added to LogRecord + return msg, kwargs + def set_json_enabled_function(self, json_enabled): """ Set a function to be called to check whether JSON output is enabled.