Skip to content

Commit 3effaee

Browse files
mdeousmdeous
committed
avoid using global variables, use a "store" class
1 parent c70f2a4 commit 3effaee

File tree

1 file changed

+55
-88
lines changed

1 file changed

+55
-88
lines changed

‎lib/utils/api.py‎

Lines changed: 55 additions & 88 deletions
Original file line numberDiff line numberDiff line change
@@ -43,17 +43,17 @@
4343
RESTAPI_SERVER_HOST = "127.0.0.1"
4444
RESTAPI_SERVER_PORT = 8775
4545

46-
# Local global variables
47-
adminid = ""
48-
db = None
49-
db_filepath = None
50-
tasks = dict()
46+
47+
# global settings
48+
class DataStore(object):
49+
admin_id = ""
50+
current_db = None
51+
tasks = dict()
5152

5253

5354
# API objects
5455
class Database(object):
55-
global db_filepath
56-
56+
filepath = None
5757
LOGS_TABLE = ("CREATE TABLE logs("
5858
"id INTEGER PRIMARY KEY AUTOINCREMENT, "
5959
"taskid INTEGER, time TEXT, "
@@ -70,10 +70,7 @@ class Database(object):
7070
")")
7171

7272
def __init__(self, database=None):
73-
if database:
74-
self.database = database
75-
else:
76-
self.database = db_filepath
73+
self.database = self.filepath if database is None else database
7774

7875
def connect(self, who="server"):
7976
self.connection = sqlite3.connect(self.database, timeout=3, isolation_level=None)
@@ -103,7 +100,6 @@ def init(self):
103100

104101

105102
class Task(object):
106-
global db_filepath
107103

108104
def __init__(self, taskid):
109105
self.process = None
@@ -125,7 +121,7 @@ def initialize_options(self, taskid):
125121
# the task ID and the file path of the IPC database
126122
self.options.api = True
127123
self.options.taskid = taskid
128-
self.options.database = db_filepath
124+
self.options.database = Database.filepath
129125

130126
# Enforce batch mode and disable coloring and ETA
131127
self.options.batch = True
@@ -269,11 +265,7 @@ def setRestAPILog():
269265

270266
# Generic functions
271267
def is_admin(taskid):
272-
global adminid
273-
if adminid != taskid:
274-
return False
275-
else:
276-
return True
268+
return DataStore.admin_id == taskid
277269

278270

279271
@hook("after_request")
@@ -328,10 +320,8 @@ def task_new():
328320
"""
329321
Create new task ID
330322
"""
331-
global tasks
332-
333323
taskid = hexencode(os.urandom(8))
334-
tasks[taskid] = Task(taskid)
324+
DataStore.tasks[taskid] = Task(taskid)
335325

336326
logger.debug("Created new task ID: %s" % taskid)
337327
return jsonize({"success": True, "taskid": taskid})
@@ -342,9 +332,9 @@ def task_delete(taskid):
342332
"""
343333
Delete own task ID
344334
"""
345-
if taskid in tasks:
346-
tasks[taskid].clean_filesystem()
347-
tasks.pop(taskid)
335+
if taskid in DataStore.tasks:
336+
DataStore.tasks[taskid].clean_filesystem()
337+
DataStore.tasks.pop(taskid)
348338

349339
logger.debug("Deleted task ID: %s" % taskid)
350340
return jsonize({"success": True})
@@ -362,8 +352,8 @@ def task_list(taskid):
362352
"""
363353
if is_admin(taskid):
364354
logger.debug("Listed task pull")
365-
task_list = list(tasks)
366-
return jsonize({"success": True, "tasks": task_list, "tasks_num": len(tasks)})
355+
task_list = list(DataStore.tasks)
356+
return jsonize({"success": True, "tasks": task_list, "tasks_num": len(task_list)})
367357
else:
368358
return jsonize({"success": False, "message": "Unauthorized"})
369359

@@ -373,13 +363,11 @@ def task_flush(taskid):
373363
"""
374364
Flush task spool (delete all tasks)
375365
"""
376-
global tasks
377-
378366
if is_admin(taskid):
379-
for task in tasks:
380-
tasks[task].clean_filesystem()
367+
for task in DataStore.tasks:
368+
DataStore.tasks[task].clean_filesystem()
381369

382-
tasks = dict()
370+
DataStore.tasks = dict()
383371
logger.debug("Flushed task pull")
384372
return jsonize({"success": True})
385373
else:
@@ -395,26 +383,24 @@ def option_list(taskid):
395383
"""
396384
List options for a certain task ID
397385
"""
398-
if taskid not in tasks:
386+
if taskid not in DataStore.tasks:
399387
return jsonize({"success": False, "message": "Invalid task ID"})
400388

401-
return jsonize({"success": True, "options": tasks[taskid].get_options()})
389+
return jsonize({"success": True, "options": DataStore.tasks[taskid].get_options()})
402390

403391

404392
@post("/option/<taskid>/get")
405393
def option_get(taskid):
406394
"""
407395
Get the value of an option (command line switch) for a certain task ID
408396
"""
409-
global tasks
410-
411-
if taskid not in tasks:
397+
if taskid not in DataStore.tasks:
412398
return jsonize({"success": False, "message": "Invalid task ID"})
413399

414400
option = request.json.get("option", "")
415401

416-
if option in tasks[taskid].options:
417-
return jsonize({"success": True, option: tasks[taskid].get_option(option)})
402+
if option in DataStore.tasks[taskid].options:
403+
return jsonize({"success": True, option: DataStore.tasks[taskid].get_option(option)})
418404
else:
419405
return jsonize({"success": False, "message": "Unknown option", option: "not set"})
420406

@@ -424,13 +410,11 @@ def option_set(taskid):
424410
"""
425411
Set an option (command line switch) for a certain task ID
426412
"""
427-
global tasks
428-
429-
if taskid not in tasks:
413+
if taskid not in DataStore.tasks:
430414
return jsonize({"success": False, "message": "Invalid task ID"})
431415

432416
for option, value in request.json.items():
433-
tasks[taskid].set_option(option, value)
417+
DataStore.tasks[taskid].set_option(option, value)
434418

435419
return jsonize({"success": True})
436420

@@ -440,36 +424,32 @@ def scan_start(taskid):
440424
"""
441425
Launch a scan
442426
"""
443-
global tasks
444-
445-
if taskid not in tasks:
427+
if taskid not in DataStore.tasks:
446428
return jsonize({"success": False, "message": "Invalid task ID"})
447429

448430
# Initialize sqlmap engine's options with user's provided options, if any
449431
for option, value in request.json.items():
450-
tasks[taskid].set_option(option, value)
432+
DataStore.tasks[taskid].set_option(option, value)
451433

452434
# Overwrite output directory value to a temporary directory
453-
tasks[taskid].set_output_directory()
435+
DataStore.tasks[taskid].set_output_directory()
454436

455437
# Launch sqlmap engine in a separate process
456-
tasks[taskid].engine_start()
438+
DataStore.tasks[taskid].engine_start()
457439

458440
logger.debug("Started scan for task ID %s" % taskid)
459-
return jsonize({"success": True, "engineid": tasks[taskid].engine_get_id()})
441+
return jsonize({"success": True, "engineid": DataStore.tasks[taskid].engine_get_id()})
460442

461443

462444
@get("/scan/<taskid>/stop")
463445
def scan_stop(taskid):
464446
"""
465447
Stop a scan
466448
"""
467-
global tasks
468-
469-
if taskid not in tasks:
449+
if taskid not in DataStore.tasks:
470450
return jsonize({"success": False, "message": "Invalid task ID"})
471451

472-
tasks[taskid].engine_stop()
452+
DataStore.tasks[taskid].engine_stop()
473453

474454
logger.debug("Stopped scan for task ID %s" % taskid)
475455
return jsonize({"success": True})
@@ -480,12 +460,10 @@ def scan_kill(taskid):
480460
"""
481461
Kill a scan
482462
"""
483-
global tasks
484-
485-
if taskid not in tasks:
463+
if taskid not in DataStore.tasks:
486464
return jsonize({"success": False, "message": "Invalid task ID"})
487465

488-
tasks[taskid].engine_kill()
466+
DataStore.tasks[taskid].engine_kill()
489467

490468
logger.debug("Killed scan for task ID %s" % taskid)
491469
return jsonize({"success": True})
@@ -496,18 +474,16 @@ def scan_status(taskid):
496474
"""
497475
Returns status of a scan
498476
"""
499-
global tasks
500-
501-
if taskid not in tasks:
477+
if taskid not in DataStore.tasks:
502478
return jsonize({"success": False, "message": "Invalid task ID"})
503479

504-
status = "terminated" if tasks[taskid].engine_has_terminated() is True else "running"
480+
status = "terminated" if DataStore.tasks[taskid].engine_has_terminated() is True else "running"
505481

506482
logger.debug("Requested status of scan for task ID %s" % taskid)
507483
return jsonize({
508484
"success": True,
509485
"status": status,
510-
"returncode": tasks[taskid].engine_get_returncode()
486+
"returncode": DataStore.tasks[taskid].engine_get_returncode()
511487
})
512488

513489

@@ -516,24 +492,23 @@ def scan_data(taskid):
516492
"""
517493
Retrieve the data of a scan
518494
"""
519-
global db
520-
global tasks
521495
json_data_message = list()
522496
json_errors_message = list()
523497

524-
if taskid not in tasks:
498+
if taskid not in DataStore.tasks:
525499
return jsonize({"success": False, "message": "Invalid task ID"})
526500

527501
# Read all data from the IPC database for the taskid
528-
for status, content_type, value in db.execute(
502+
for status, content_type, value in DataStore.current_db.execute(
529503
"SELECT status, content_type, value FROM data WHERE taskid = ? ORDER BY id ASC",
530504
(taskid,)):
531505
json_data_message.append(
532506
{"status": status, "type": content_type, "value": dejsonize(value)})
533507

534508
# Read all error messages from the IPC database
535-
for error in db.execute("SELECT error FROM errors WHERE taskid = ? ORDER BY id ASC",
536-
(taskid,)):
509+
for error in DataStore.current_db.execute(
510+
"SELECT error FROM errors WHERE taskid = ? ORDER BY id ASC",
511+
(taskid,)):
537512
json_errors_message.append(error)
538513

539514
logger.debug("Retrieved data and error messages for scan for task ID %s" % taskid)
@@ -545,11 +520,9 @@ def scan_log_limited(taskid, start, end):
545520
"""
546521
Retrieve a subset of log messages
547522
"""
548-
global db
549-
global tasks
550523
json_log_messages = list()
551524

552-
if taskid not in tasks:
525+
if taskid not in DataStore.tasks:
553526
return jsonize({"success": False, "message": "Invalid task ID"})
554527

555528
if not start.isdigit() or not end.isdigit() or end < start:
@@ -559,7 +532,7 @@ def scan_log_limited(taskid, start, end):
559532
end = max(1, int(end))
560533

561534
# Read a subset of log messages from the IPC database
562-
for time_, level, message in db.execute(
535+
for time_, level, message in DataStore.current_db.execute(
563536
("SELECT time, level, message FROM logs WHERE "
564537
"taskid = ? AND id >= ? AND id <= ? ORDER BY id ASC"),
565538
(taskid, start, end)):
@@ -574,15 +547,13 @@ def scan_log(taskid):
574547
"""
575548
Retrieve the log messages
576549
"""
577-
global db
578-
global tasks
579550
json_log_messages = list()
580551

581-
if taskid not in tasks:
552+
if taskid not in DataStore.tasks:
582553
return jsonize({"success": False, "message": "Invalid task ID"})
583554

584555
# Read all log messages from the IPC database
585-
for time_, level, message in db.execute(
556+
for time_, level, message in DataStore.current_db.execute(
586557
"SELECT time, level, message FROM logs WHERE taskid = ? ORDER BY id ASC", (taskid,)):
587558
json_log_messages.append({"time": time_, "level": level, "message": message})
588559

@@ -595,7 +566,7 @@ def download(taskid, target, filename):
595566
"""
596567
Download a certain file from the file system
597568
"""
598-
if taskid not in tasks:
569+
if taskid not in DataStore.tasks:
599570
return jsonize({"success": False, "message": "Invalid task ID"})
600571

601572
# Prevent file path traversal - the lame way
@@ -616,21 +587,17 @@ def server(host="0.0.0.0", port=RESTAPI_SERVER_PORT):
616587
"""
617588
REST-JSON API server
618589
"""
619-
global adminid
620-
global db
621-
global db_filepath
622-
623-
adminid = hexencode(os.urandom(16))
624-
db_filepath = tempfile.mkstemp(prefix="sqlmapipc-", text=False)[1]
590+
DataStore.admin_id = hexencode(os.urandom(16))
591+
Database.filepath = tempfile.mkstemp(prefix="sqlmapipc-", text=False)[1]
625592

626593
logger.info("Running REST-JSON API server at '%s:%d'.." % (host, port))
627-
logger.info("Admin ID: %s" % adminid)
628-
logger.debug("IPC database: %s" % db_filepath)
594+
logger.info("Admin ID: %s" % DataStore.admin_id)
595+
logger.debug("IPC database: %s" % Database.filepath)
629596

630597
# Initialize IPC database
631-
db = Database()
632-
db.connect()
633-
db.init()
598+
DataStore.current_db = Database()
599+
DataStore.current_db.connect()
600+
DataStore.current_db.init()
634601

635602
# Run RESTful API
636603
run(host=host, port=port, quiet=True, debug=False)

0 commit comments

Comments
 (0)