def announce_failed_tasks(self, event): """ Called when a celery task fails. Is used to reset dates and scanner of the targeted tool. Args: event: created automatically when the event occurs. Contains some info about the task """ self.state.event(event) task = self.state.tasks.get(event['uuid']) # get task # Get tasks arguments try: argsWere = json.loads(task.info()["args"].replace("'", "\"")) args = argsWere[0] toolId = ObjectId(args[1]) # args[1] is the tool_id for task_running in self.tasks_running: if str(task_running[1]) == str(toolId): task_running[0].revoke(terminate=True) del task_running tool = Tool.fetchObject({"_id": ObjectId(toolId)}) tool.markAsNotDone() except InvalidId: pass except json.decoder.JSONDecodeError: pass except KeyError: pass # Plugin failed so "args" does not exist
def findLaunchableTools(): """ Try to find tools that matches all criteria. Returns: A tuple with two values: * A list of launchable tools as dictionary with values _id, name and priority * A dictionary of waiting tools with tool's names as keys and integer as value. """ toolsLaunchable = [] waiting = {} time_compatible_waves_id = Wave.searchForAddressCompatibleWithTime() for wave_id in time_compatible_waves_id: commandsLaunchableWave = Wave.getNotDoneTools(wave_id) for tool in commandsLaunchableWave: toolModel = Tool.fetchObject({"_id": tool}) try: waiting[str(toolModel)] += 1 except KeyError: waiting[str(toolModel)] = 1 command = toolModel.getCommand() if command is None: prio = 0 else: prio = int(command.get("priority", 0)) toolsLaunchable.append({ "_id": tool, "name": str(toolModel), "priority": prio }) return toolsLaunchable, waiting
def dispatchLaunchableTools(my_monitor, launchableTools): """ Try to launch given tools within the monitor Args: my_monitor: A Monitor instance which knows what tools are already launched and online workers launchableTools: A list of tools within a Wave that passed the Intervals checking. """ mongoInstance = MongoCalendar.getInstance() for launchableTool in launchableTools: tool = Tool.fetchObject({"_id": ObjectId(launchableTool["_id"])}) my_monitor.launchTask(mongoInstance.calendarName, tool)
def findLaunchableToolsOnWorker(worker, calendarName): """ Try to find tools that matches all criteria. Args: workerName: the current working worker Returns: A tuple with two values: * A list of launchable tools as dictionary with values _id, name and priority * A dictionary of waiting tools with tool's names as keys and integer as value. """ mongoInstance = MongoCalendar.getInstance() mongoInstance.connectToDb(calendarName) toolsLaunchable = [] worker_registered = mongoInstance.findInDb("pollenisator", "workers", {"name": worker.name}, False) commands_registered = worker_registered["registeredCommands"] waiting = {} time_compatible_waves_id = Wave.searchForAddressCompatibleWithTime() for wave_id in time_compatible_waves_id: commandsLaunchableWave = Wave.getNotDoneTools(wave_id) for tool in commandsLaunchableWave: toolModel = Tool.fetchObject({"_id": tool}) if toolModel.name not in commands_registered: continue if worker.hasRegistered(toolModel): try: waiting[str(toolModel)] += 1 except KeyError: waiting[str(toolModel)] = 1 command = toolModel.getCommand() if command is None: prio = 0 else: prio = int(command.get("priority", 0)) toolsLaunchable.append({ "_id": tool, "name": str(toolModel), "priority": prio, "errored": "error" in toolModel.status }) return toolsLaunchable, waiting
def executeCommand(calendarName, toolId, parser=""): """ CELERY remote task Execute the tool with the given toolId on the given calendar name. Then execute the plugin corresponding. Any unhandled exception will result in a task-failed event in the Monitor class. Args: calendarName: The calendar to search the given tool id for. toolId: the mongo Object id corresponding to the tool to execute. parser: plugin name to execute. If empty, the plugin specified in tools.d will be feteched. Raises: Terminated: if the task gets terminated OSError: if the output directory cannot be created (not if it already exists) Exception: if an exception unhandled occurs during the bash command execution. Exception: if a plugin considered a failure. """ # Connect to given calendar mongoInstance = MongoCalendar.getInstance() mongoInstance.connectToDb(calendarName) msg = "" # retrieve tool from tool sid toolModel = Tool.fetchObject({"_id": ObjectId(toolId)}) if toolModel is None: raise Exception("Tool does not exist : " + str(toolId)) command = Command.fetchObject({"name": toolModel.name}, calendarName) # Get time limit and output directory if toolModel.wave == "Custom commands": timeLimit = None else: timeLimit = getWaveTimeLimit(toolModel.wave) if command is not None: timeLimit = min(datetime.now() + timedelta(0, int(command.timeout)), timeLimit) outputRelDir = toolModel.getOutputDir(calendarName) abs_path = os.path.dirname(os.path.abspath(__file__)) outputDir = os.path.join(abs_path, "./results", outputRelDir) # Create the output directory try: os.makedirs(outputDir) except OSError as exc: if exc.errno == errno.EEXIST and os.path.isdir(outputDir): pass else: raise exc # Read Tool config file tools_infos = Utils.loadToolsConfig() comm = toolModel.getCommandToExecute(outputDir) if parser.strip() == "": if toolModel.name not in list(tools_infos.keys()): msg = "TASK FAILED Received tool that was not registered : " + \ str(toolModel.name)+" not in "+str(list(tools_infos.keys())) raise Exception(msg) # Fetch the command to execute if tools_infos.get(toolModel.name, None) is None: bin_path = "" else: bin_path = tools_infos[toolModel.name].get("bin") if bin_path is not None: if not bin_path.endswith(" "): bin_path = bin_path + " " comm = bin_path + comm if comm != "": try: # Load the plugin if parser.strip() == "": mod = Utils.loadPlugin(tools_infos[toolModel.name]["plugin"]) elif parser.strip() == "auto-detect": mod = Utils.loadPluginByBin(toolModel.name.split("::")[0]) else: mod = Utils.loadPlugin(parser) # Complete command with file output toolFileName = toolModel.name+"_" + \ str(time.time())+mod.getFileOutputExt() comm = mod.changeCommand(comm, outputDir, toolFileName) print(('TASK STARTED:' + toolModel.name)) print("Will timeout at " + str(timeLimit)) # Execute the command with a timeout returncode = Utils.execute(comm, timeLimit, True) except Exception as e: raise e # Execute found plugin if there is one if mod is not None: filepath = mod.getFileOutputPath(comm) try: # Open generated file as utf8 with io.open(filepath, "r", encoding="utf-8", errors='ignore') as file_opened: # Check return code by plugin (can be always true if the return code is inconsistent) if mod.checkReturnCode(returncode): notes, tags, _, _ = mod.Parse(file_opened) if notes is None: notes = "No results found by plugin." if tags is None: tags = [] if isinstance(tags, str): tags = [tags] # Success could be change to False by the plugin function (evaluating the return code for exemple) # if the success is validated, mark tool as done toolModel.markAsDone( os.path.join(outputRelDir, os.path.basename(filepath))) # And update the tool in database toolModel.notes = notes toolModel.tags = tags toolModel.update() # Upload file to SFTP mod.centralizeFile(filepath, outputDir) msg = "TASK SUCCESS : " + toolModel.name else: # BAS RESULT OF PLUGIN msg = "TASK FAILED (says the mod) : " + toolModel.name msg += "The return code was not the expected one. (" + str( returncode) + ")." toolModel.markAsError() raise Exception(msg) except IOError as e: toolModel.tags = ["todo"] toolModel.notes = "Failed to read results file" toolModel.markAsDone() else: msg = "TASK FAILED (no plugin found) : " + toolModel.name toolModel.markAsNotDone() raise Exception(msg) # Delay if command is not None: if float(command.sleep_between) > 0.0: msg += " (will sleep for " + \ str(float(command.sleep_between))+")" print(msg) time.sleep(float(command.sleep_between)) return
def doExecuteCommand(workerToken, calendarName, toolId): """ remote task Execute the tool with the given toolId on the given calendar name. Then execute the plugin corresponding. Any unhandled exception will result in a task-failed event in the class. Args: calendarName: The calendar to search the given tool id for. toolId: the mongo Object id corresponding to the tool to execute. parser: plugin name to execute. If empty, the plugin specified in tools.d will be feteched. Raises: Terminated: if the task gets terminated OSError: if the output directory cannot be created (not if it already exists) Exception: if an exception unhandled occurs during the bash command execution. Exception: if a plugin considered a failure. """ apiclient = APIClient.getInstance() apiclient.setToken(workerToken) apiclient.currentPentest = calendarName # bypass login by not using connectToDb toolModel = Tool.fetchObject({"_id":ObjectId(toolId)}) command_dict = toolModel.getCommand() msg = "" success, comm, fileext = apiclient.getCommandline(toolId) if not success: print(str(comm)) toolModel.setStatus(["error"]) return False, str(comm) bin_path = command_dict["bin_path"] if bin_path is not None: if not bin_path.endswith(" "): bin_path = bin_path+" " comm = bin_path+comm outputRelDir = toolModel.getOutputDir(calendarName) abs_path = os.path.dirname(os.path.abspath(__file__)) toolFileName = toolModel.name+"_" + \ str(time.time()) # ext already added in command outputDir = os.path.join(abs_path, "./results", outputRelDir) # Create the output directory try: os.makedirs(outputDir) except OSError as exc: if exc.errno == errno.EEXIST and os.path.isdir(outputDir): pass else: print(str(exc)) toolModel.setStatus(["error"]) return False, str(exc) outputDir = os.path.join(outputDir, toolFileName) comm = comm.replace("|outputDir|", outputDir) # Get tool's wave time limit searching the wave intervals if toolModel.wave == "Custom commands": timeLimit = None else: timeLimit = getWaveTimeLimit(toolModel.wave) # adjust timeLimit if the command has a lower timeout if command_dict is not None: timeLimit = min(datetime.now()+timedelta(0, int(command_dict.get("timeout", 0))), timeLimit) ## if "timedout" in toolModel.status: timeLimit = None try: toolModel.text = comm toolModel.update({"text":comm}) print(('TASK STARTED:'+toolModel.name)) print("Will timeout at "+str(timeLimit)) # Execute the command with a timeout returncode = Utils.execute(comm, timeLimit, True) if returncode == -1: toolModel.setStatus(["timedout"]) return False, str("Command timedout") except Exception as e: print(str(e)) toolModel.setStatus(["error"]) return False, str(e) # Execute found plugin if there is one outputfile = outputDir+fileext plugin = "auto-detect" if command_dict["plugin"] == "" else command_dict["plugin"] msg = apiclient.importToolResult(toolId, plugin, outputfile) if msg != "Success": #toolModel.markAsNotDone() print(str(msg)) toolModel.setStatus(["error"]) return False, str(msg) # Delay if command_dict is not None: if float(command_dict.get("sleep_between", 0)) > 0.0: msg += " (will sleep for " + \ str(float(command_dict.get("sleep_between", 0)))+")" print(msg) time.sleep(float(command_dict.get("sleep_between", 0))) return True, ""