From c3f20a136f4ae569d801eafab24873c7af232615 Mon Sep 17 00:00:00 2001 From: Miroslav Stampar Date: Wed, 12 Dec 2012 14:03:03 +0100 Subject: [PATCH] Minor update for an Issue #287 --- _sqlmap.py | 24 ++++++++++++++++++++--- lib/controller/controller.py | 7 +++++++ lib/core/log.py | 4 ++++ lib/core/option.py | 1 + lib/utils/xmlrpc.py | 38 ++++++++++++++++++++++++------------ 5 files changed, 58 insertions(+), 16 deletions(-) diff --git a/_sqlmap.py b/_sqlmap.py index a18ccb9bd..30144bfd6 100755 --- a/_sqlmap.py +++ b/_sqlmap.py @@ -8,6 +8,7 @@ See the file 'doc/COPYING' for copying permission import bdb import logging import os +import StringIO import sys import time import traceback @@ -22,6 +23,7 @@ from lib.core.common import dataToStdout from lib.core.common import getUnicode from lib.core.common import setPaths from lib.core.common import weAreFrozen +from lib.core.convert import stdoutencode from lib.core.data import cmdLineOptions from lib.core.data import conf from lib.core.data import kb @@ -31,6 +33,8 @@ from lib.core.common import unhandledExceptionMessage from lib.core.exception import exceptionsTuple from lib.core.exception import SqlmapSilentQuitException from lib.core.exception import SqlmapUserQuitException +from lib.core.log import FORMATTER +from lib.core.log import setLoggerHandler from lib.core.option import init from lib.core.profiling import profile from lib.core.settings import LEGAL_DISCLAIMER @@ -48,6 +52,22 @@ def modulePath(): return os.path.dirname(getUnicode(sys.executable if weAreFrozen() else __file__, sys.getfilesystemencoding())) +def xmlRpcServe(): + logger.setLevel(logging.INFO) + server = XMLRPCServer(cmdLineOptions.xmlRpcPort or XMLRPC_SERVER_PORT) + class _(logging.Handler): + def emit(self, record): + message = stdoutencode(self.format(record)) + sys.stdout.write("%s\n" % message) + handler = _() + handler.is_tty = False + handler.disableColoring = True + handler.setFormatter(FORMATTER) + setLoggerHandler(handler) + sys.stdout = StringIO.StringIO() + sys.stderr = StringIO.StringIO() + server.serve() + def main(): """ Main function of sqlmap when running from command line. @@ -65,9 +85,7 @@ def main(): cmdLineOptions.update(cmdLineParser().__dict__) if cmdLineOptions.xmlRpc: - logger.setLevel(logging.INFO) - server = XMLRPCServer(cmdLineOptions.xmlRpcPort or XMLRPC_SERVER_PORT) - server.serve() + xmlRpcServe() else: init(cmdLineOptions) diff --git a/lib/controller/controller.py b/lib/controller/controller.py index 70221bae5..12626d016 100644 --- a/lib/controller/controller.py +++ b/lib/controller/controller.py @@ -228,6 +228,13 @@ def _saveToResultsFile(): conf.resultsFP.writelines(line) def start(): + kb.busyFlag = True + retVal = _start() + kb.busyFlag = False + + return retVal + +def _start(): """ This function calls a function that performs checks on both URL stability and all GET, POST, Cookie and User-Agent parameters to diff --git a/lib/core/log.py b/lib/core/log.py index 091b287bb..72c825998 100644 --- a/lib/core/log.py +++ b/lib/core/log.py @@ -31,3 +31,7 @@ FORMATTER = logging.Formatter("\r[%(asctime)s] [%(levelname)s] %(message)s", "%H LOGGER_HANDLER.setFormatter(FORMATTER) LOGGER.addHandler(LOGGER_HANDLER) LOGGER.setLevel(logging.WARN) + +def setLoggerHandler(handler): + LOGGER.handlers = [] + LOGGER.addHandler(handler) diff --git a/lib/core/option.py b/lib/core/option.py index da2328428..3ca56b9eb 100644 --- a/lib/core/option.py +++ b/lib/core/option.py @@ -1487,6 +1487,7 @@ def _setKnowledgeBaseAttributes(flushAll=True): kb.brute = AttribDict({"tables":[], "columns":[]}) kb.bruteMode = False + kb.busyFlag = False kb.cache = AttribDict() kb.cache.content = {} diff --git a/lib/utils/xmlrpc.py b/lib/utils/xmlrpc.py index 29ea8a9e0..b0c3c0db2 100644 --- a/lib/utils/xmlrpc.py +++ b/lib/utils/xmlrpc.py @@ -6,6 +6,7 @@ See the file 'doc/COPYING' for copying permission """ import sys +import threading import xmlrpclib try: @@ -14,6 +15,7 @@ try: from lib.controller.controller import start from lib.core.datatype import AttribDict from lib.core.data import cmdLineOptions + from lib.core.data import kb from lib.core.data import logger from lib.core.option import init from lib.core.settings import UNICODE_ENCODING @@ -26,6 +28,16 @@ class XMLRPCServer: self.port = port self.reset() + self.server = SimpleXMLRPCServer(addr=("", self.port), logRequests=False, allow_none=True, encoding=UNICODE_ENCODING) + self.server.register_function(self.reset) + self.server.register_function(self.set_option) + self.server.register_function(self.get_option) + self.server.register_function(self.get_option_names) + self.server.register_function(self.run) + logger.info("Registering RPC methods: %s" % str(self.server.system_listMethods()).strip("[]")) + self.server.register_introspection_functions() + logger.info("Running XML-RPC server at '0.0.0.0:%d'..." % self.port) + def reset(self): self.options = AttribDict(cmdLineOptions) @@ -38,21 +50,20 @@ class XMLRPCServer: def get_option_names(self): return self.options.keys() + def is_busy(self): + return kb.get("busyFlag") + def run(self): - init(self.options, True) - return start() + if not self.is_busy(): + init(self.options, True) + thread = threading.Thread(target=start) + thread.daemon = True + thread.start() + else: + raise Exception, "sqlmap busy" def serve(self): - server = SimpleXMLRPCServer(addr=("", self.port), logRequests=False, allow_none=True, encoding=UNICODE_ENCODING) - server.register_function(self.reset) - server.register_function(self.set_option) - server.register_function(self.get_option) - server.register_function(self.get_option_names) - server.register_function(self.run) - logger.info("Registering RPC methods: %s" % str(server.system_listMethods()).strip("[]")) - server.register_introspection_functions() - logger.info("Running XML-RPC server at '0.0.0.0:%d'..." % self.port) - server.serve_forever() + self.server.serve_forever() if __name__ == "__main__": try: @@ -69,7 +80,8 @@ if __name__ == "__main__": print "[i] Server instance name: 'server'" print "[i] Sample usage: 'server.system.listMethods()'" except Exception, ex: - print "[x] '%s'" % str(ex) + if ex: + print "[x] '%s'" % str(ex) else: while True: try: