Skip to content

Commit 8a94650

Browse files
mdeousmdeous
committed
PEP8
1 parent 5b2ded0 commit 8a94650

1 file changed

Lines changed: 65 additions & 19 deletions

File tree

‎lib/utils/api.py‎

Lines changed: 65 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
#!/usr/bin/env python
2+
# -*- coding: utf-8 -*-
23

34
"""
45
Copyright (c) 2006-2013 sqlmap developers (http://sqlmap.org/)
@@ -52,13 +53,25 @@
5253
db_filepath = None
5354
tasks = dict()
5455

56+
5557
# API objects
5658
class Database(object):
5759
global db_filepath
5860

59-
LOGS_TABLE = "CREATE TABLE logs(id INTEGER PRIMARY KEY AUTOINCREMENT, taskid INTEGER, time TEXT, level TEXT, message TEXT)"
60-
DATA_TABLE = "CREATE TABLE data(id INTEGER PRIMARY KEY AUTOINCREMENT, taskid INTEGER, status INTEGER, content_type INTEGER, value TEXT)"
61-
ERRORS_TABLE = "CREATE TABLE errors(id INTEGER PRIMARY KEY AUTOINCREMENT, taskid INTEGER, error TEXT)"
61+
LOGS_TABLE = ("CREATE TABLE logs("
62+
"id INTEGER PRIMARY KEY AUTOINCREMENT, "
63+
"taskid INTEGER, time TEXT, "
64+
"level TEXT, message TEXT"
65+
")")
66+
DATA_TABLE = ("CREATE TABLE data("
67+
"id INTEGER PRIMARY KEY AUTOINCREMENT, "
68+
"taskid INTEGER, status INTEGER, "
69+
"content_type INTEGER, value TEXT"
70+
")")
71+
ERRORS_TABLE = ("CREATE TABLE errors("
72+
"id INTEGER PRIMARY KEY AUTOINCREMENT, "
73+
"taskid INTEGER, error TEXT"
74+
")")
6275

6376
def __init__(self, database=None):
6477
if database:
@@ -92,6 +105,7 @@ def init(self):
92105
self.execute(self.DATA_TABLE)
93106
self.execute(self.ERRORS_TABLE)
94107

108+
95109
class Task(object):
96110
global db_filepath
97111

@@ -111,7 +125,8 @@ def initialize_options(self, taskid):
111125
type_ = unArrayizeValue(type_)
112126
self.options[name] = _defaults.get(name, datatype[type_])
113127

114-
# Let sqlmap engine knows it is getting called by the API, the task ID and the file path of the IPC database
128+
# Let sqlmap engine knows it is getting called by the API,
129+
# the task ID and the file path of the IPC database
115130
self.options.api = True
116131
self.options.taskid = taskid
117132
self.options.database = db_filepath
@@ -145,7 +160,8 @@ def clean_filesystem(self):
145160
shutil.rmtree(self.output_directory)
146161

147162
def engine_start(self):
148-
self.process = Popen("python sqlmap.py --pickled-options %s" % base64pickle(self.options), shell=True, stdin=PIPE, close_fds=False)
163+
self.process = Popen("python sqlmap.py --pickled-options %s" % base64pickle(self.options),
164+
shell=True, stdin=PIPE, close_fds=False)
149165

150166
def engine_stop(self):
151167
if self.process:
@@ -194,25 +210,27 @@ def write(self, value, status=CONTENT_STATUS.IN_PROGRESS, content_type=None):
194210
# Ignore all non-relevant messages
195211
return
196212

197-
output = conf.database_cursor.execute("SELECT id, status, value FROM data WHERE taskid = ? AND content_type = ?",
198-
(self.taskid, content_type))
199-
200-
#print >>sys.__stdout__, "output: %s\nvalue: %s\nstatus: %d\ncontent_type: %d\nkb.partRun: %s\n--------------" % (output, value, status, content_type, kb.partRun)
213+
output = conf.database_cursor.execute(
214+
"SELECT id, status, value FROM data WHERE taskid = ? AND content_type = ?",
215+
(self.taskid, content_type))
201216

202217
# Delete partial output from IPC database if we have got a complete output
203218
if status == CONTENT_STATUS.COMPLETE:
204219
if len(output) > 0:
205220
for index in xrange(0, len(output)):
206-
conf.database_cursor.execute("DELETE FROM data WHERE id = ?", (output[index][0],))
221+
conf.database_cursor.execute("DELETE FROM data WHERE id = ?",
222+
(output[index][0],))
207223

208-
conf.database_cursor.execute("INSERT INTO data VALUES(NULL, ?, ?, ?, ?)", (self.taskid, status, content_type, jsonize(value)))
224+
conf.database_cursor.execute("INSERT INTO data VALUES(NULL, ?, ?, ?, ?)",
225+
(self.taskid, status, content_type, jsonize(value)))
209226
if kb.partRun:
210227
kb.partRun = None
211228

212229
elif status == CONTENT_STATUS.IN_PROGRESS:
213230
if len(output) == 0:
214231
conf.database_cursor.execute("INSERT INTO data VALUES(NULL, ?, ?, ?, ?)",
215-
(self.taskid, status, content_type, jsonize(value)))
232+
(self.taskid, status, content_type,
233+
jsonize(value)))
216234
else:
217235
new_value = "%s%s" % (dejsonize(output[0][2]), value)
218236
conf.database_cursor.execute("UPDATE data SET value = ? WHERE id = ?",
@@ -230,6 +248,7 @@ def close(self):
230248
def seek(self):
231249
pass
232250

251+
233252
class LogRecorder(logging.StreamHandler):
234253
def emit(self, record):
235254
"""
@@ -238,7 +257,8 @@ def emit(self, record):
238257
"""
239258
conf.database_cursor.execute("INSERT INTO logs VALUES(NULL, ?, ?, ?, ?)",
240259
(conf.taskid, time.strftime("%X"), record.levelname,
241-
record.msg % record.args if record.args else record.msg))
260+
record.msg % record.args if record.args else record.msg))
261+
242262

243263
def setRestAPILog():
244264
if hasattr(conf, "api"):
@@ -250,6 +270,7 @@ def setRestAPILog():
250270
LOGGER_RECORDER = LogRecorder()
251271
logger.addHandler(LOGGER_RECORDER)
252272

273+
253274
# Generic functions
254275
def is_admin(taskid):
255276
global adminid
@@ -258,6 +279,7 @@ def is_admin(taskid):
258279
else:
259280
return True
260281

282+
261283
@hook("after_request")
262284
def security_headers(json_header=True):
263285
"""
@@ -282,16 +304,19 @@ def error401(error=None):
282304
security_headers(False)
283305
return "Access denied"
284306

307+
285308
@error(404) # Not Found
286309
def error404(error=None):
287310
security_headers(False)
288311
return "Nothing here"
289312

313+
290314
@error(405) # Method Not Allowed (e.g. when requesting a POST method via GET)
291315
def error405(error=None):
292316
security_headers(False)
293317
return "Method not allowed"
294318

319+
295320
@error(500) # Internal Server Error
296321
def error500(error=None):
297322
security_headers(False)
@@ -315,6 +340,7 @@ def task_new():
315340
logger.debug("Created new task ID: %s" % taskid)
316341
return jsonize({"taskid": taskid})
317342

343+
318344
@get("/task/<taskid>/delete")
319345
def task_delete(taskid):
320346
"""
@@ -345,6 +371,7 @@ def task_list(taskid):
345371
else:
346372
abort(401)
347373

374+
348375
@get("/admin/<taskid>/flush")
349376
def task_flush(taskid):
350377
"""
@@ -377,6 +404,7 @@ def option_list(taskid):
377404

378405
return jsonize({"options": tasks[taskid].get_options()})
379406

407+
380408
@post("/option/<taskid>/get")
381409
def option_get(taskid):
382410
"""
@@ -394,6 +422,7 @@ def option_get(taskid):
394422
else:
395423
return jsonize({option: "not set"})
396424

425+
397426
@post("/option/<taskid>/set")
398427
def option_set(taskid):
399428
"""
@@ -435,6 +464,7 @@ def scan_start(taskid):
435464
logger.debug("Started scan for task ID %s" % taskid)
436465
return jsonize({"success": True, "engineid": tasks[taskid].engine_get_id()})
437466

467+
438468
@get("/scan/<taskid>/stop")
439469
def scan_stop(taskid):
440470
"""
@@ -450,6 +480,7 @@ def scan_stop(taskid):
450480
logger.debug("Stopped scan for task ID %s" % taskid)
451481
return jsonize({"success": True})
452482

483+
453484
@get("/scan/<taskid>/kill")
454485
def scan_kill(taskid):
455486
"""
@@ -465,6 +496,7 @@ def scan_kill(taskid):
465496
logger.debug("Killed scan for task ID %s" % taskid)
466497
return jsonize({"success": True})
467498

499+
468500
@get("/scan/<taskid>/status")
469501
def scan_status(taskid):
470502
"""
@@ -480,6 +512,7 @@ def scan_status(taskid):
480512
logger.debug("Requested status of scan for task ID %s" % taskid)
481513
return jsonize({"status": status, "returncode": tasks[taskid].engine_get_returncode()})
482514

515+
483516
@get("/scan/<taskid>/data")
484517
def scan_data(taskid):
485518
"""
@@ -494,11 +527,15 @@ def scan_data(taskid):
494527
abort(500, "Invalid task ID")
495528

496529
# Read all data from the IPC database for the taskid
497-
for status, content_type, value in db.execute("SELECT status, content_type, value FROM data WHERE taskid = ? ORDER BY id ASC", (taskid,)):
498-
json_data_message.append({"status": status, "type": content_type, "value": dejsonize(value)})
530+
for status, content_type, value in db.execute(
531+
"SELECT status, content_type, value FROM data WHERE taskid = ? ORDER BY id ASC",
532+
(taskid,)):
533+
json_data_message.append(
534+
{"status": status, "type": content_type, "value": dejsonize(value)})
499535

500536
# Read all error messages from the IPC database
501-
for error in db.execute("SELECT error FROM errors WHERE taskid = ? ORDER BY id ASC", (taskid,)):
537+
for error in db.execute("SELECT error FROM errors WHERE taskid = ? ORDER BY id ASC",
538+
(taskid,)):
502539
json_errors_message.append(error)
503540

504541
logger.debug("Retrieved data and error messages for scan for task ID %s" % taskid)
@@ -524,12 +561,16 @@ def scan_log_limited(taskid, start, end):
524561
end = max(1, int(end))
525562

526563
# Read a subset of log messages from the IPC database
527-
for time_, level, message in db.execute("SELECT time, level, message FROM logs WHERE taskid = ? AND id >= ? AND id <= ? ORDER BY id ASC", (taskid, start, end)):
564+
for time_, level, message in db.execute(
565+
("SELECT time, level, message FROM logs WHERE "
566+
"taskid = ? AND id >= ? AND id <= ? ORDER BY id ASC"),
567+
(taskid, start, end)):
528568
json_log_messages.append({"time": time_, "level": level, "message": message})
529569

530570
logger.debug("Retrieved subset of log messages for scan for task ID %s" % taskid)
531571
return jsonize({"log": json_log_messages})
532572

573+
533574
@get("/scan/<taskid>/log")
534575
def scan_log(taskid):
535576
"""
@@ -543,7 +584,8 @@ def scan_log(taskid):
543584
abort(500, "Invalid task ID")
544585

545586
# Read all log messages from the IPC database
546-
for time_, level, message in db.execute("SELECT time, level, message FROM logs WHERE taskid = ? ORDER BY id ASC", (taskid,)):
587+
for time_, level, message in db.execute(
588+
"SELECT time, level, message FROM logs WHERE taskid = ? ORDER BY id ASC", (taskid,)):
547589
json_log_messages.append({"time": time_, "level": level, "message": message})
548590

549591
logger.debug("Retrieved log messages for scan for task ID %s" % taskid)
@@ -569,6 +611,7 @@ def download(taskid, target, filename):
569611
else:
570612
abort(500, "File does not exist")
571613

614+
572615
def server(host="0.0.0.0", port=RESTAPI_SERVER_PORT):
573616
"""
574617
REST-JSON API server
@@ -592,6 +635,7 @@ def server(host="0.0.0.0", port=RESTAPI_SERVER_PORT):
592635
# Run RESTful API
593636
run(host=host, port=port, quiet=True, debug=False)
594637

638+
595639
def client(host=RESTAPI_SERVER_HOST, port=RESTAPI_SERVER_PORT):
596640
"""
597641
REST-JSON API client
@@ -602,6 +646,8 @@ def client(host=RESTAPI_SERVER_HOST, port=RESTAPI_SERVER_PORT):
602646
# TODO: write a simple client with requests, for now use curl from command line
603647
logger.error("Not yet implemented, use curl from command line instead for now, for example:")
604648
print "\n\t$ curl http://%s:%d/task/new" % (host, port)
605-
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)
649+
print ("\t$ curl -H \"Content-Type: application/json\" "
650+
"-X POST -d '{\"url\": \"http://testphp.vulnweb.com/artists.php?artist=1\"}' "
651+
"http://%s:%d/scan/:taskid/start") % (host, port)
606652
print "\t$ curl http://%s:%d/scan/:taskid/data" % (host, port)
607653
print "\t$ curl http://%s:%d/scan/:taskid/log\n" % (host, port)

0 commit comments

Comments
 (0)