From 72378d4f61f12d141c14dc78333dae2bf52829ef Mon Sep 17 00:00:00 2001
From: Miroslav Stampar <miroslav.stampar@gmail.com>
Date: Wed, 11 Jul 2012 20:29:48 +0200
Subject: [PATCH] Some more refactoring

---
 lib/core/common.py | 11 ++++++-----
 lib/core/log.py    | 20 --------------------
 2 files changed, 6 insertions(+), 25 deletions(-)

diff --git a/lib/core/common.py b/lib/core/common.py
index cdc4c0e6f..adfa4072f 100644
--- a/lib/core/common.py
+++ b/lib/core/common.py
@@ -78,9 +78,7 @@ from lib.core.exception import sqlmapMissingDependence
 from lib.core.exception import sqlmapSilentQuitException
 from lib.core.exception import sqlmapSyntaxException
 from lib.core.log import FORMATTER
-from lib.core.log import LEVEL_ATTRS
-from lib.core.log import LEVEL_COLORS
-from lib.core.log import LEVEL_ON_COLORS
+from lib.core.log import LOGGER_HANDLER
 from lib.core.optiondict import optDict
 from lib.core.settings import CUSTOM_INJECTION_MARK_CHAR
 from lib.core.settings import DEFAULT_COOKIE_DELIMITER
@@ -698,10 +696,13 @@ def singleTimeLogMessage(message, level=logging.INFO, flag=None):
         logger.log(level, message)
 
 def setColor(message, bold=False):
+    message = "[TRAFFIC IN] " + message
     level = extractRegexResult(r"\A\s*\[(?P<result>[A-Z ]+)\]", message)
 
-    if level:
-        retVal = colored(message, color=LEVEL_COLORS.get(level), on_color=LEVEL_ON_COLORS.get(level), attrs=LEVEL_ATTRS.get(level))
+    _ = LOGGER_HANDLER.level_map.get(logging._levelNames.get(level))
+    if _:
+        background, foreground, bold = _
+        retVal = colored(message, color=foreground, on_color="on_%s" % background if background else None, attrs=("bold",) if bold else None)
     else:
         retVal = message
 
diff --git a/lib/core/log.py b/lib/core/log.py
index 8467b2a0c..1e3a883a7 100644
--- a/lib/core/log.py
+++ b/lib/core/log.py
@@ -18,26 +18,6 @@ logging.addLevelName(CUSTOM_LOGGING.TRAFFIC_IN, "TRAFFIC IN")
 
 LOGGER = logging.getLogger("sqlmapLog")
 
-LEVEL_COLORS = {
-                    "CRITICAL": "white",
-                    "ERROR": "red",
-                    "WARNING": "yellow",
-                    "INFO": "green",
-                    "DEBUG": "blue",
-                    "PAYLOAD": "magenta",
-                    "TRAFFIC OUT": "cyan",
-                    "TRAFFIC IN": "grey"
-               }
-
-LEVEL_ON_COLORS = {
-                    "CRITICAL": "on_red",
-                  }
-
-LEVEL_ATTRS = {
-                "CRITICAL": ('bold',),
-              }
-
-
 try:
     import ctypes
     LOGGER_HANDLER = ColorizingStreamHandler(sys.stdout)