Minor refactoring

This commit is contained in:
Miroslav Stampar 2012-07-11 19:54:21 +02:00
parent 53ccd09ca4
commit d7926b8aac
3 changed files with 22 additions and 24 deletions

View File

@ -79,6 +79,7 @@ from lib.core.exception import sqlmapSilentQuitException
from lib.core.exception import sqlmapSyntaxException from lib.core.exception import sqlmapSyntaxException
from lib.core.log import FORMATTER from lib.core.log import FORMATTER
from lib.core.optiondict import optDict from lib.core.optiondict import optDict
from lib.core.settings import COLOR_MAP
from lib.core.settings import CUSTOM_INJECTION_MARK_CHAR from lib.core.settings import CUSTOM_INJECTION_MARK_CHAR
from lib.core.settings import DEFAULT_COOKIE_DELIMITER from lib.core.settings import DEFAULT_COOKIE_DELIMITER
from lib.core.settings import DEFAULT_GET_POST_DELIMITER from lib.core.settings import DEFAULT_GET_POST_DELIMITER
@ -694,30 +695,17 @@ def singleTimeLogMessage(message, level=logging.INFO, flag=None):
kb.singleLogFlags.add(flag) kb.singleLogFlags.add(flag)
logger.log(level, message) logger.log(level, message)
def setCurrentMessage(message): def setColor(message, bold=False):
logMsg = re.search("(CRITICAL|ERROR|WARNING|INFO|DEBUG|PAYLOAD|TRAFFIC OUT|TRAFFIC IN)", message) retVal = message
level = extractRegexResult(r"\A\s*\[(?P<result>[A-Z ]+)\]", message)
if logMsg: if level:
kb.currentMessage = logMsg.group(0) attrs = ('bold',) if (level == "CRITICAL" or bold) else None
on_color = 'on_red' if level == "CRITICAL" else None
color = COLOR_MAP.get(level)
retVal = colored(message, color=color, on_color=on_color, attrs=attrs)
def setColour(message, bold=False): return retVal
if not hasattr(kb, "currentMessage"):
return message
setCurrentMessage(message)
color_map = {
"CRITICAL": "white",
"ERROR": "red",
"WARNING": "yellow",
"INFO": "green",
"DEBUG": "blue",
"PAYLOAD": "magenta",
"TRAFFIC OUT": "cyan",
"TRAFFIC IN": "grey"
}
return colored(message, color_map[kb.currentMessage] if kb.currentMessage else None, on_color='on_red' if kb.currentMessage == "CRITICAL" else None, attrs=['bold'] if (kb.currentMessage == "CRITICAL" or bold) else None)
def dataToStdout(data, forceOutput=False, bold=False): def dataToStdout(data, forceOutput=False, bold=False):
""" """
@ -751,7 +739,7 @@ def dataToStdout(data, forceOutput=False, bold=False):
except: except:
message = data.encode(UNICODE_ENCODING) message = data.encode(UNICODE_ENCODING)
sys.stdout.write(setColour(message, bold)) sys.stdout.write(setColor(message, bold))
sys.stdout.flush() sys.stdout.flush()
if kb.get("multiThreadMode"): if kb.get("multiThreadMode"):

View File

@ -1436,7 +1436,6 @@ def __setKnowledgeBaseAttributes(flushAll=True):
kb.commonOutputs = None kb.commonOutputs = None
kb.counters = {} kb.counters = {}
kb.currentMessage = None
kb.data = AttribDict() kb.data = AttribDict()
kb.dataOutputFlag = False kb.dataOutputFlag = False

View File

@ -268,6 +268,17 @@ BASIC_HELP_ITEMS = (
"wizard" "wizard"
) )
COLOR_MAP = {
"CRITICAL": "white",
"ERROR": "red",
"WARNING": "yellow",
"INFO": "green",
"DEBUG": "blue",
"PAYLOAD": "magenta",
"TRAFFIC OUT": "cyan",
"TRAFFIC IN": "grey"
}
# string representation for NULL value # string representation for NULL value
NULL = "NULL" NULL = "NULL"