#
# Copyright (C) 2010-2012 Vinay Sajip. All rights reserved. Licensed under the new BSD license.
# (Note: 2018 modifications by @stamparm)
#
import logging
import os
import re
import sys
from lib.core.settings import IS_WIN
if IS_WIN:
import ctypes
import ctypes.wintypes
# Reference: https://gist.github.com/vsajip/758430
# https://github.com/ipython/ipython/issues/4252
# https://msdn.microsoft.com/en-us/library/windows/desktop/ms686047%28v=vs.85%29.aspx
ctypes.windll.kernel32.SetConsoleTextAttribute.argtypes = [ctypes.wintypes.HANDLE, ctypes.wintypes.WORD]
ctypes.windll.kernel32.SetConsoleTextAttribute.restype = ctypes.wintypes.BOOL
def stdoutEncode(data): # Cross-referenced function
return data
class ColorizingStreamHandler(logging.StreamHandler):
# color names to indices
color_map = {
'black': 0,
'red': 1,
'green': 2,
'yellow': 3,
'blue': 4,
'magenta': 5,
'cyan': 6,
'white': 7,
}
# levels to (background, foreground, bold/intense)
level_map = {
logging.DEBUG: (None, 'blue', False),
logging.INFO: (None, 'green', False),
logging.WARNING: (None, 'yellow', False),
logging.ERROR: (None, 'red', False),
logging.CRITICAL: ('red', 'white', False)
}
csi = '\x1b['
reset = '\x1b[0m'
bold = "\x1b[1m"
disable_coloring = False
@property
def is_tty(self):
isatty = getattr(self.stream, 'isatty', None)
return isatty and isatty() and not self.disable_coloring
def emit(self, record):
try:
message = stdoutEncode(self.format(record))
stream = self.stream
if not self.is_tty:
if message and message[0] == "\r":
message = message[1:]
stream.write(message)
else:
self.output_colorized(message)
stream.write(getattr(self, 'terminator', '\n'))
self.flush()
except (KeyboardInterrupt, SystemExit):
raise
except IOError:
pass
except:
self.handleError(record)
if not IS_WIN:
def output_colorized(self, message):
self.stream.write(message)
else:
ansi_esc = re.compile(r'\x1b\[((?:\d+)(?:;(?:\d+))*)m')
nt_color_map = {
0: 0x00, # black
1: 0x04, # red
2: 0x02, # green
3: 0x06, # yellow
4: 0x01, # blue
5: 0x05, # magenta
6: 0x03, # cyan
7: 0x07, # white
}
def output_colorized(self, message):
parts = self.ansi_esc.split(message)
h = None
fd = getattr(self.stream, 'fileno', None)
if fd is not None:
fd = fd()
if fd in (1, 2): # stdout or stderr
h = ctypes.windll.kernel32.GetStdHandle(-10 - fd)
while parts:
text = parts.pop(0)
if text:
self.stream.write(text)
self.stream.flush()
if parts:
params = parts.pop(0)
if h is not None:
params = [int(p) for p in params.split(';')]
color = 0
for p in params:
if 40 <= p <= 47:
color |= self.nt_color_map[p - 40] << 4
elif 30 <= p <= 37:
color |= self.nt_color_map[p - 30]
elif p == 1:
color |= 0x08 # foreground intensity on
elif p == 0: # reset to default color
color = 0x07
else:
pass # error condition ignored
ctypes.windll.kernel32.SetConsoleTextAttribute(h, color)
def _reset(self, message):
if not message.endswith(self.reset):
reset = self.reset
elif self.bold in message: # bold
reset = self.reset + self.bold
else:
reset = self.reset
return reset
def colorize(self, message, levelno):
if levelno in self.level_map and self.is_tty:
bg, fg, bold = self.level_map[levelno]
params = []
if bg in self.color_map:
params.append(str(self.color_map[bg] + 40))
if fg in self.color_map:
params.append(str(self.color_map[fg] + 30))
if bold:
params.append('1')
if params and message:
match = re.search(r"\A(\s+)", message)
prefix = match.group(1) if match else ""
message = message[len(prefix):]
match = re.search(r"\[([A-Z ]+)\]", message) # log level
if match:
level = match.group(1)
if message.startswith(self.bold):
message = message.replace(self.bold, "")
reset = self.reset + self.bold
params.append('1')
else:
reset = self.reset
message = message.replace(level, ''.join((self.csi, ';'.join(params), 'm', level, reset)), 1)
match = re.search(r"\A\s*\[([\d:]+)\]", message) # time
if match:
time = match.group(1)
message = message.replace(time, ''.join((self.csi, str(self.color_map["cyan"] + 30), 'm', time, self._reset(message))), 1)
match = re.search(r"\[(#\d+)\]", message) # counter
if match:
counter = match.group(1)
message = message.replace(counter, ''.join((self.csi, str(self.color_map["yellow"] + 30), 'm', counter, self._reset(message))), 1)
if level != "PAYLOAD":
if any(_ in message for _ in ("parsed DBMS error message",)):
match = re.search(r": '(.+)'", message)
if match:
string = match.group(1)
message = message.replace("'%s'" % string, "'%s'" % ''.join((self.csi, str(self.color_map["white"] + 30), 'm', string, self._reset(message))), 1)
else:
match = re.search(r"\bresumed: '(.+\.\.\.)", message)
if match:
string = match.group(1)
message = message.replace("'%s" % string, "'%s" % ''.join((self.csi, str(self.color_map["white"] + 30), 'm', string, self._reset(message))), 1)
else:
match = re.search(r" \('(.+)'\)\Z", message)
if match:
string = match.group(1)
message = message.replace("'%s'" % string, "'%s'" % ''.join((self.csi, str(self.color_map["white"] + 30), 'm', string, self._reset(message))), 1)
else:
for match in re.finditer(r"[^\w]'([^']+)'", message): # single-quoted
string = match.group(1)
message = message.replace("'%s'" % string, "'%s'" % ''.join((self.csi, str(self.color_map["white"] + 30), 'm', string, self._reset(message))), 1)
else:
message = ''.join((self.csi, ';'.join(params), 'm', message, self.reset))
if prefix:
message = "%s%s" % (prefix, message)
message = message.replace("%s]" % self.bold, "]%s" % self.bold) # dirty patch
return message
def format(self, record):
message = logging.StreamHandler.format(self, record)
return self.colorize(message, record.levelno)