From 1152cf8958eee13e827503812f3902f01d7a9c8f Mon Sep 17 00:00:00 2001 From: Bernardo Damele Date: Tue, 29 Jan 2013 15:38:09 +0000 Subject: [PATCH] increased SQLite connection timeout to 3 seconds, the object will now wait for the lock to go away max 3 seconds, no longer 1 only. Relevant code refactoring and minor improvements all over the API library (issue #297) --- lib/utils/api.py | 120 +++++++++++++++++++++++------------------------ 1 file changed, 59 insertions(+), 61 deletions(-) diff --git a/lib/utils/api.py b/lib/utils/api.py index 926673536..ca640d255 100644 --- a/lib/utils/api.py +++ b/lib/utils/api.py @@ -60,10 +60,10 @@ class Database(object): def create(self): _, self.database = tempfile.mkstemp(prefix="sqlmapipc-", text=False) - logger.info("IPC database is %s" % self.database) + logger.debug("IPC database: %s" % self.database) def connect(self): - self.connection = sqlite3.connect(self.database, timeout=1, isolation_level=None) + self.connection = sqlite3.connect(self.database, timeout=3, isolation_level=None) self.cursor = self.connection.cursor() def disconnect(self): @@ -132,18 +132,28 @@ class Task(object): shutil.rmtree(self.output_directory) def engine_start(self): - self.process = Popen("python sqlmap.py --pickled-options %s" % base64pickle(self.options), shell=True, stdin=PIPE) + self.process = Popen("python sqlmap.py --pickled-options %s" % base64pickle(self.options), shell=True, stdin=PIPE, close_fds=False) def engine_stop(self): if self.process: - self.process.terminate() + return self.process.terminate() + else: + return None def engine_kill(self): if self.process: - self.process.kill() + return self.process.kill() + else: + return None - def engine_get_pid(self): - return self.processid.pid + def engine_get_id(self): + if self.process: + return self.process.pid + else: + return None + + def engine_has_terminated(self): + return isinstance(self.process.returncode, int) == True # Wrapper functions for sqlmap engine class StdDbOut(object): @@ -162,9 +172,13 @@ class StdDbOut(object): def write(self, value, status=None, content_type=None): if self.messagetype == "stdout": - conf.database_cursor.execute("INSERT INTO data VALUES(NULL, ?, ?, ?, ?)", (self.taskid, status, content_type, jsonize(value))) + #conf.database_cursor.execute("INSERT INTO data VALUES(NULL, ?, ?, ?, ?)", + # (self.taskid, status, content_type, base64pickle(value))) + conf.database_cursor.execute("INSERT INTO data VALUES(NULL, ?, ?, ?, ?)", + (self.taskid, status, content_type, jsonize(value))) else: - conf.database_cursor.execute("INSERT INTO errors VALUES(NULL, ?, ?)", (self.taskid, value)) + conf.database_cursor.execute("INSERT INTO errors VALUES(NULL, ?, ?)", + (self.taskid, str(value) if value else "")) def flush(self): pass @@ -182,8 +196,7 @@ class LogRecorder(logging.StreamHandler): communication with the parent process """ conf.database_cursor.execute("INSERT INTO logs VALUES(NULL, ?, ?, ?, ?)", - (conf.taskid, time.strftime("%X"), record.levelname, - record.msg % record.args if record.args else record.msg)) + (conf.taskid, time.strftime("%X"), record.levelname, record.msg % record.args if record.args else record.msg)) def setRestAPILog(): if hasattr(conf, "api"): @@ -257,35 +270,44 @@ def task_new(): taskid = hexencode(os.urandom(8)) tasks[taskid] = Task(taskid) + logger.debug("Created new task ID: %s" % taskid) + return jsonize({"taskid": taskid}) -@get("/task//destroy") -def task_destroy(taskid): +@get("/task//delete") +def task_delete(taskid): """ - Destroy own task ID + Delete own task ID """ if taskid in tasks: tasks[taskid].clean_filesystem() tasks.pop(taskid) + + logger.debug("Deleted task ID: %s" % taskid) + return jsonize({"success": True}) else: abort(500, "Invalid task ID") -# Admin's methods -@get("/task//list") +################### +# Admin functions # +################### + +@get("/admin//list") def task_list(taskid): """ - List all active tasks + List task poll """ if is_admin(taskid): - return jsonize({"tasks": tasks}) + logger.debug("Listed task poll") + return jsonize({"tasks": tasks, "tasks_num": len(tasks)}) else: abort(401) -@get("/task//flush") +@get("/admin//flush") def task_flush(taskid): """ - Flush task spool (destroy all tasks) + Flush task spool (delete all tasks) """ global tasks @@ -294,6 +316,7 @@ def task_flush(taskid): tasks[task].clean_filesystem() tasks = dict() + logger.debug("Flushed task poll") return jsonize({"success": True}) else: abort(401) @@ -302,20 +325,7 @@ def task_flush(taskid): # sqlmap core interact functions # ################################## -# Admin's methods -@get("/status/") -def status(taskid): - """ - Verify the status of the API as well as the core - """ - - if is_admin(taskid): - tasks_num = len(tasks) - return jsonize({"tasks": tasks_num}) - else: - abort(401) - -# Functions to handle options +# Handle task's options @get("/option//list") def option_list(taskid): """ @@ -324,7 +334,7 @@ def option_list(taskid): if taskid not in tasks: abort(500, "Invalid task ID") - return jsonize(tasks[taskid].get_options()) + return jsonize({"options": tasks[taskid].get_options()}) @post("/option//get") def option_get(taskid): @@ -339,7 +349,7 @@ def option_get(taskid): if option in tasks[taskid]: return jsonize({option: tasks[taskid].get_option(option)}) else: - return jsonize({option: None}) + return jsonize({option: "Not set"}) @post("/option//set") def option_set(taskid): @@ -356,7 +366,7 @@ def option_set(taskid): return jsonize({"success": True}) -# Function to handle scans +# Handle scans @post("/scan//start") def scan_start(taskid): """ @@ -375,12 +385,12 @@ def scan_start(taskid): tasks[taskid].set_output_directory() # Launch sqlmap engine in a separate thread - logger.debug("starting a scan for task ID %s" % taskid) + logger.debug("Starting a scan for task ID %s" % taskid) # Launch sqlmap engine tasks[taskid].engine_start() - return jsonize({"success": True}) + return jsonize({"success": True, "engineid": tasks[taskid].engine_get_id()}) @get("/scan//stop") def scan_stop(taskid): @@ -406,21 +416,6 @@ def scan_kill(taskid): return jsonize({"success": tasks[taskid].engine_kill()}) -@get("/scan//delete") -def scan_delete(taskid): - """ - Delete a scan and corresponding temporary output directory and IPC database - """ - global tasks - - if taskid not in tasks: - abort(500, "Invalid task ID") - - scan_stop(taskid) - tasks[taskid].clean_filesystem() - - return jsonize({"success": True}) - @get("/scan//data") def scan_data(taskid): """ @@ -436,7 +431,8 @@ def scan_data(taskid): # Read all data from the IPC database for the taskid for status, content_type, value in db.execute("SELECT status, content_type, value FROM data WHERE taskid = ? ORDER BY id ASC", (taskid,)): - json_data_message.append([status, content_type, dejsonize(value)]) + #json_data_message.append({"status": status, "type": content_type, "value": base64unpickle(value)}) + json_data_message.append({"status": status, "type": content_type, "value": dejsonize(value)}) # Read all error messages from the IPC database for error in db.execute("SELECT error FROM errors WHERE taskid = ? ORDER BY id ASC", (taskid,)): @@ -515,24 +511,26 @@ def server(host="0.0.0.0", port=RESTAPI_SERVER_PORT): global db adminid = hexencode(os.urandom(16)) + + logger.info("Running REST-JSON API server at '%s:%d'.." % (host, port)) + logger.info("Admin ID: %s" % adminid) + + # Initialize IPC database db = Database() db.initialize() - logger.info("running REST-JSON API server at '%s:%d'.." % (host, port)) - logger.info("the admin task ID is: %s" % adminid) - # Run RESTful API - run(host=host, port=port, quiet=False, debug=False) + run(host=host, port=port, quiet=True, debug=False) def client(host=RESTAPI_SERVER_HOST, port=RESTAPI_SERVER_PORT): """ REST-JSON API client """ addr = "http://%s:%d" % (host, port) - logger.info("starting debug REST-JSON client to '%s'..." % addr) + logger.info("Starting REST-JSON API client to '%s'..." % addr) # TODO: write a simple client with requests, for now use curl from command line - logger.error("not yet implemented, use curl from command line instead for now, for example:") + logger.error("Not yet implemented, use curl from command line instead for now, for example:") print "\n\t$ curl http://%s:%d/task/new" % (host, port) print "\t$ curl -H \"Content-Type: application/json\" -X POST -d '{\"url\": \"http://testphp.vulnweb.com/artists.php?artist=1\"}' http://%s:%d/scan/:taskid/start" % (host, port) print "\t$ curl http://%s:%d/scan/:taskid/output" % (host, port)