Thanks to visit codestin.com
Credit goes to github.com

Skip to content

Commit 1152cf8

Browse files
bdamelebdamele
committed
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)
1 parent 9677e0f commit 1152cf8

1 file changed

Lines changed: 59 additions & 61 deletions

File tree

‎lib/utils/api.py‎

Lines changed: 59 additions & 61 deletions
Original file line numberDiff line numberDiff line change
@@ -60,10 +60,10 @@ def __init__(self):
6060

6161
def create(self):
6262
_, self.database = tempfile.mkstemp(prefix="sqlmapipc-", text=False)
63-
logger.info("IPC database is %s" % self.database)
63+
logger.debug("IPC database: %s" % self.database)
6464

6565
def connect(self):
66-
self.connection = sqlite3.connect(self.database, timeout=1, isolation_level=None)
66+
self.connection = sqlite3.connect(self.database, timeout=3, isolation_level=None)
6767
self.cursor = self.connection.cursor()
6868

6969
def disconnect(self):
@@ -132,18 +132,28 @@ def clean_filesystem(self):
132132
shutil.rmtree(self.output_directory)
133133

134134
def engine_start(self):
135-
self.process = Popen("python sqlmap.py --pickled-options %s" % base64pickle(self.options), shell=True, stdin=PIPE)
135+
self.process = Popen("python sqlmap.py --pickled-options %s" % base64pickle(self.options), shell=True, stdin=PIPE, close_fds=False)
136136

137137
def engine_stop(self):
138138
if self.process:
139-
self.process.terminate()
139+
return self.process.terminate()
140+
else:
141+
return None
140142

141143
def engine_kill(self):
142144
if self.process:
143-
self.process.kill()
145+
return self.process.kill()
146+
else:
147+
return None
148+
149+
def engine_get_id(self):
150+
if self.process:
151+
return self.process.pid
152+
else:
153+
return None
144154

145-
def engine_get_pid(self):
146-
return self.processid.pid
155+
def engine_has_terminated(self):
156+
return isinstance(self.process.returncode, int) == True
147157

148158
# Wrapper functions for sqlmap engine
149159
class StdDbOut(object):
@@ -162,9 +172,13 @@ def __init__(self, taskid, messagetype="stdout"):
162172

163173
def write(self, value, status=None, content_type=None):
164174
if self.messagetype == "stdout":
165-
conf.database_cursor.execute("INSERT INTO data VALUES(NULL, ?, ?, ?, ?)", (self.taskid, status, content_type, jsonize(value)))
175+
#conf.database_cursor.execute("INSERT INTO data VALUES(NULL, ?, ?, ?, ?)",
176+
# (self.taskid, status, content_type, base64pickle(value)))
177+
conf.database_cursor.execute("INSERT INTO data VALUES(NULL, ?, ?, ?, ?)",
178+
(self.taskid, status, content_type, jsonize(value)))
166179
else:
167-
conf.database_cursor.execute("INSERT INTO errors VALUES(NULL, ?, ?)", (self.taskid, value))
180+
conf.database_cursor.execute("INSERT INTO errors VALUES(NULL, ?, ?)",
181+
(self.taskid, str(value) if value else ""))
168182

169183
def flush(self):
170184
pass
@@ -182,8 +196,7 @@ def emit(self, record):
182196
communication with the parent process
183197
"""
184198
conf.database_cursor.execute("INSERT INTO logs VALUES(NULL, ?, ?, ?, ?)",
185-
(conf.taskid, time.strftime("%X"), record.levelname,
186-
record.msg % record.args if record.args else record.msg))
199+
(conf.taskid, time.strftime("%X"), record.levelname, record.msg % record.args if record.args else record.msg))
187200

188201
def setRestAPILog():
189202
if hasattr(conf, "api"):
@@ -257,35 +270,44 @@ def task_new():
257270
taskid = hexencode(os.urandom(8))
258271
tasks[taskid] = Task(taskid)
259272

273+
logger.debug("Created new task ID: %s" % taskid)
274+
260275
return jsonize({"taskid": taskid})
261276

262-
@get("/task/<taskid>/destroy")
263-
def task_destroy(taskid):
277+
@get("/task/<taskid>/delete")
278+
def task_delete(taskid):
264279
"""
265-
Destroy own task ID
280+
Delete own task ID
266281
"""
267282
if taskid in tasks:
268283
tasks[taskid].clean_filesystem()
269284
tasks.pop(taskid)
285+
286+
logger.debug("Deleted task ID: %s" % taskid)
287+
270288
return jsonize({"success": True})
271289
else:
272290
abort(500, "Invalid task ID")
273291

274-
# Admin's methods
275-
@get("/task/<taskid>/list")
292+
###################
293+
# Admin functions #
294+
###################
295+
296+
@get("/admin/<taskid>/list")
276297
def task_list(taskid):
277298
"""
278-
List all active tasks
299+
List task poll
279300
"""
280301
if is_admin(taskid):
281-
return jsonize({"tasks": tasks})
302+
logger.debug("Listed task poll")
303+
return jsonize({"tasks": tasks, "tasks_num": len(tasks)})
282304
else:
283305
abort(401)
284306

285-
@get("/task/<taskid>/flush")
307+
@get("/admin/<taskid>/flush")
286308
def task_flush(taskid):
287309
"""
288-
Flush task spool (destroy all tasks)
310+
Flush task spool (delete all tasks)
289311
"""
290312
global tasks
291313

@@ -294,6 +316,7 @@ def task_flush(taskid):
294316
tasks[task].clean_filesystem()
295317

296318
tasks = dict()
319+
logger.debug("Flushed task poll")
297320
return jsonize({"success": True})
298321
else:
299322
abort(401)
@@ -302,20 +325,7 @@ def task_flush(taskid):
302325
# sqlmap core interact functions #
303326
##################################
304327

305-
# Admin's methods
306-
@get("/status/<taskid>")
307-
def status(taskid):
308-
"""
309-
Verify the status of the API as well as the core
310-
"""
311-
312-
if is_admin(taskid):
313-
tasks_num = len(tasks)
314-
return jsonize({"tasks": tasks_num})
315-
else:
316-
abort(401)
317-
318-
# Functions to handle options
328+
# Handle task's options
319329
@get("/option/<taskid>/list")
320330
def option_list(taskid):
321331
"""
@@ -324,7 +334,7 @@ def option_list(taskid):
324334
if taskid not in tasks:
325335
abort(500, "Invalid task ID")
326336

327-
return jsonize(tasks[taskid].get_options())
337+
return jsonize({"options": tasks[taskid].get_options()})
328338

329339
@post("/option/<taskid>/get")
330340
def option_get(taskid):
@@ -339,7 +349,7 @@ def option_get(taskid):
339349
if option in tasks[taskid]:
340350
return jsonize({option: tasks[taskid].get_option(option)})
341351
else:
342-
return jsonize({option: None})
352+
return jsonize({option: "Not set"})
343353

344354
@post("/option/<taskid>/set")
345355
def option_set(taskid):
@@ -356,7 +366,7 @@ def option_set(taskid):
356366

357367
return jsonize({"success": True})
358368

359-
# Function to handle scans
369+
# Handle scans
360370
@post("/scan/<taskid>/start")
361371
def scan_start(taskid):
362372
"""
@@ -375,12 +385,12 @@ def scan_start(taskid):
375385
tasks[taskid].set_output_directory()
376386

377387
# Launch sqlmap engine in a separate thread
378-
logger.debug("starting a scan for task ID %s" % taskid)
388+
logger.debug("Starting a scan for task ID %s" % taskid)
379389

380390
# Launch sqlmap engine
381391
tasks[taskid].engine_start()
382392

383-
return jsonize({"success": True})
393+
return jsonize({"success": True, "engineid": tasks[taskid].engine_get_id()})
384394

385395
@get("/scan/<taskid>/stop")
386396
def scan_stop(taskid):
@@ -406,21 +416,6 @@ def scan_kill(taskid):
406416

407417
return jsonize({"success": tasks[taskid].engine_kill()})
408418

409-
@get("/scan/<taskid>/delete")
410-
def scan_delete(taskid):
411-
"""
412-
Delete a scan and corresponding temporary output directory and IPC database
413-
"""
414-
global tasks
415-
416-
if taskid not in tasks:
417-
abort(500, "Invalid task ID")
418-
419-
scan_stop(taskid)
420-
tasks[taskid].clean_filesystem()
421-
422-
return jsonize({"success": True})
423-
424419
@get("/scan/<taskid>/data")
425420
def scan_data(taskid):
426421
"""
@@ -436,7 +431,8 @@ def scan_data(taskid):
436431

437432
# Read all data from the IPC database for the taskid
438433
for status, content_type, value in db.execute("SELECT status, content_type, value FROM data WHERE taskid = ? ORDER BY id ASC", (taskid,)):
439-
json_data_message.append([status, content_type, dejsonize(value)])
434+
#json_data_message.append({"status": status, "type": content_type, "value": base64unpickle(value)})
435+
json_data_message.append({"status": status, "type": content_type, "value": dejsonize(value)})
440436

441437
# Read all error messages from the IPC database
442438
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):
515511
global db
516512

517513
adminid = hexencode(os.urandom(16))
514+
515+
logger.info("Running REST-JSON API server at '%s:%d'.." % (host, port))
516+
logger.info("Admin ID: %s" % adminid)
517+
518+
# Initialize IPC database
518519
db = Database()
519520
db.initialize()
520521

521-
logger.info("running REST-JSON API server at '%s:%d'.." % (host, port))
522-
logger.info("the admin task ID is: %s" % adminid)
523-
524522
# Run RESTful API
525-
run(host=host, port=port, quiet=False, debug=False)
523+
run(host=host, port=port, quiet=True, debug=False)
526524

527525
def client(host=RESTAPI_SERVER_HOST, port=RESTAPI_SERVER_PORT):
528526
"""
529527
REST-JSON API client
530528
"""
531529
addr = "http://%s:%d" % (host, port)
532-
logger.info("starting debug REST-JSON client to '%s'..." % addr)
530+
logger.info("Starting REST-JSON API client to '%s'..." % addr)
533531

534532
# TODO: write a simple client with requests, for now use curl from command line
535-
logger.error("not yet implemented, use curl from command line instead for now, for example:")
533+
logger.error("Not yet implemented, use curl from command line instead for now, for example:")
536534
print "\n\t$ curl http://%s:%d/task/new" % (host, port)
537535
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)
538536
print "\t$ curl http://%s:%d/scan/:taskid/output" % (host, port)

0 commit comments

Comments
 (0)