From 1417decdf1c7b284cd714bcb3bfe26a1da14839f Mon Sep 17 00:00:00 2001 From: daremon Date: Mon, 14 Sep 2015 17:31:02 +0300 Subject: [PATCH 1/2] Added commands stop, kill, list to API client --- lib/utils/api.py | 33 ++++++++++++++++++++++++--------- 1 file changed, 24 insertions(+), 9 deletions(-) diff --git a/lib/utils/api.py b/lib/utils/api.py index 45eb46be8..d66097261 100644 --- a/lib/utils/api.py +++ b/lib/utils/api.py @@ -368,18 +368,19 @@ def task_delete(taskid): ################### +@get("/admin/list") @get("/admin//list") -def task_list(taskid): +def task_list(taskid=None): """ List task pull """ - if is_admin(taskid): - logger.debug("[%s] Listed task pool" % taskid) + logger.debug("[%s] Listed task pool") + if taskid is not None: tasks = list(DataStore.tasks) - return jsonize({"success": True, "tasks": tasks, "tasks_num": len(tasks)}) else: - logger.warning("[%s] Unauthorized call to task_list()" % taskid) - return jsonize({"success": False, "message": "Unauthorized"}) + tasks = {x: dejsonize(scan_status(x))['status'] + for x in list(DataStore.tasks)} + return jsonize({"success": True, "tasks": tasks, "tasks_num": len(tasks)}) @get("/admin//flush") @@ -475,7 +476,9 @@ def scan_stop(taskid): """ Stop a scan """ - if taskid not in DataStore.tasks: + if (taskid not in DataStore.tasks or + DataStore.tasks[taskid].engine_process() is None or + DataStore.tasks[taskid].engine_has_terminated()): logger.warning("[%s] Invalid task ID provided to scan_stop()" % taskid) return jsonize({"success": False, "message": "Invalid task ID"}) @@ -490,7 +493,9 @@ def scan_kill(taskid): """ Kill a scan """ - if taskid not in DataStore.tasks: + if (taskid not in DataStore.tasks or + DataStore.tasks[taskid].engine_process() is None or + DataStore.tasks[taskid].engine_has_terminated()): logger.warning("[%s] Invalid task ID provided to scan_kill()" % taskid) return jsonize({"success": False, "message": "Invalid task ID"}) @@ -691,7 +696,7 @@ def client(host=RESTAPI_SERVER_HOST, port=RESTAPI_SERVER_PORT): print break - if command.lower() in ("data", "log", "status"): + if command.lower() in ("data", "log", "status", "stop", "kill"): if not taskid: logger.error("No task ID in use") continue @@ -743,6 +748,13 @@ def client(host=RESTAPI_SERVER_HOST, port=RESTAPI_SERVER_PORT): continue logger.info("Switching to task ID '%s' " % taskid) + elif command.lower() == "list": + raw = _client(addr + "/admin/list") + res = dejsonize(raw) + if not res["success"]: + logger.error("Failed to execute command " + command) + dataToStdout("%s\n" % raw) + elif command.lower() in ("exit", "bye", "quit", 'q'): return @@ -753,6 +765,9 @@ def client(host=RESTAPI_SERVER_HOST, port=RESTAPI_SERVER_PORT): msg += "data Retrieve and show data for current task\n" msg += "log Retrieve and show log for current task\n" msg += "status Retrieve and show status for current task\n" + msg += "stop Stop current task\n" + msg += "kill Kill current task\n" + msg += "list Display all tasks\n" msg += "exit Exit this client\n" dataToStdout(msg) From c2fb2161d34e21f6cece0e3d5069c4c9c78bdd76 Mon Sep 17 00:00:00 2001 From: daremon Date: Wed, 16 Sep 2015 00:15:16 +0300 Subject: [PATCH 2/2] Added flush command --- lib/utils/api.py | 17 +++++++++-------- 1 file changed, 9 insertions(+), 8 deletions(-) diff --git a/lib/utils/api.py b/lib/utils/api.py index 7a73905a8..09aac2c5a 100644 --- a/lib/utils/api.py +++ b/lib/utils/api.py @@ -698,12 +698,12 @@ def client(host=RESTAPI_SERVER_HOST, port=RESTAPI_SERVER_PORT): while True: try: - command = raw_input("api%s> " % (" (%s)" % taskid if taskid else "")).strip() + command = raw_input("api%s> " % (" (%s)" % taskid if taskid else "")).strip().lower() except (EOFError, KeyboardInterrupt): print break - if command.lower() in ("data", "log", "status", "stop", "kill"): + if command in ("data", "log", "status", "stop", "kill"): if not taskid: logger.error("No task ID in use") continue @@ -713,7 +713,7 @@ def client(host=RESTAPI_SERVER_HOST, port=RESTAPI_SERVER_PORT): logger.error("Failed to execute command " + command) dataToStdout("%s\n" % raw) - elif command.lower().startswith("new"): + elif command.startswith("new"): if ' ' not in command: logger.error("Program arguments are missing") continue @@ -745,7 +745,7 @@ def client(host=RESTAPI_SERVER_HOST, port=RESTAPI_SERVER_PORT): continue logger.info("Scanning started") - elif command.lower().startswith("use"): + elif command.startswith("use"): taskid = (command.split()[1] if ' ' in command else "").strip("'\"") if not taskid: logger.error("Task ID is missing") @@ -757,17 +757,17 @@ def client(host=RESTAPI_SERVER_HOST, port=RESTAPI_SERVER_PORT): continue logger.info("Switching to task ID '%s' " % taskid) - elif command.lower() == "list": - raw = _client(addr + "/admin/0/list") + elif command in ("list", "flush"): + raw = _client(addr + "/admin/0/" + command) res = dejsonize(raw) if not res["success"]: logger.error("Failed to execute command " + command) dataToStdout("%s\n" % raw) - elif command.lower() in ("exit", "bye", "quit", 'q'): + elif command in ("exit", "bye", "quit", 'q'): return - elif command.lower() in ("help", "?"): + elif command in ("help", "?"): msg = "help Show this help message\n" msg += "new ARGS Start a new scan task with provided arguments (e.g. 'new -u \"http://testphp.vulnweb.com/artists.php?artist=1\"')\n" msg += "use TASKID Switch current context to different task (e.g. 'use c04d8c5c7582efb4')\n" @@ -777,6 +777,7 @@ def client(host=RESTAPI_SERVER_HOST, port=RESTAPI_SERVER_PORT): msg += "stop Stop current task\n" msg += "kill Kill current task\n" msg += "list Display all tasks\n" + msg += "flush Flush tasks (delete all tasks)\n" msg += "exit Exit this client\n" dataToStdout(msg)