def addOnlineWorker(self, worker_hostname):
"""
Register a celery worker on the worker's list. Also deletes old queues and messages
Args:
worker_hostname: the worker hostname to register on worker's list
"""
mongoInstance = MongoCalendar.getInstance()
agg_queues = mongoInstance.aggregateFromDb(
"broker_pollenisator", "messages.routing", [{
"$group": {
"_id": "$queue"
}
}, {
"$match": {
"_id": {
"$regex": "^.*&" + worker_hostname + "&.*$"
}
}
}])
mongoInstance.deleteFromDb(
"broker_pollenisator", "messages.routing",
{"queue": {
"$regex": "^.*&" + worker_hostname + "&.*$"
}}, True)
for agg_queue in agg_queues:
Utils.execute(
"celery -A slave purge -f -Q '" + agg_queue["_id"] + "'", None,
False)
self.workerRegisterCommands(worker_hostname)
def open(self):
mainDir = os.path.normpath(Utils.getMainDir())
mongoInstance = MongoCalendar.getInstance()
with open(os.path.join(mainDir, "setupTerminalForPentest.sh"), "r") as f:
data = f.read()
lines = data.split("\n")
lines[0] = "POLLENISATOR_CURRENT_DB="+str(mongoInstance.calendarName)
data = "\n".join(lines)
with open(os.path.join(mainDir, "setupTerminalForPentest.sh"), "w") as f:
f.write(data)
favorite = self.settings.getFavoriteTerm()
if favorite is None:
tkinter.messagebox.showerror("Terminal settings invalid", "None of the terminals given in the settings are installed on this computer.")
return False
if which(favorite) is not None:
terms = self.settings.getTerms()
terms_dict = {}
for term in terms:
terms_dict[term.split(" ")[0]] = term
command_term = terms_dict.get(favorite, None)
if command_term is not None:
Utils.execute(terms_dict[favorite])
return True
else:
tkinter.messagebox.showerror("Terminal settings invalid", "Check your terminal settings")
else:
tkinter.messagebox.showerror("Terminal settings invalid", "The selected favorite terminal is not available on this computer.")
return False
def downloadResultFile(self, _event=None):
"""Callback for tool click #TODO move to ToolController
Download the tool result file and asks the user if he or she wants to open it.
If OK, tries to open it using xdg-open or os.startsfile
Args:
_event: not used
"""
fs = FileStorage()
fs.open()
path = None
if fs.sftp_connection is not None:
dialog = ChildDialogInfo(self.appliViewFrame, "Download Started",
"Downloading...")
resultFile = self.controller.getResultFile()
dialog.show()
if resultFile != "" and resultFile is not None:
path = fs.getToolResult(resultFile)
else:
tkinter.messagebox.showerror(
"Download failed", "The result file does not exist.")
dialog.destroy()
else:
tkinter.messagebox.showerror("Download failed",
"The sftp connection failed.")
return
fs.close()
if path is not None:
if os.path.isfile(path):
if which("xdg-open") is not None:
dialog = ChildDialogQuestion(
self.appliViewFrame,
"Download completed",
"The file has been downloaded.\n Would you like to open it?",
answers=["Open", "Cancel"])
self.appliViewFrame.wait_window(dialog.app)
if dialog.rvalue == "Open":
Utils.execute("xdg-open " + path)
return
else:
return
path = None
if path is None:
tkinter.messagebox.showerror(
"Download failed", "the file does not exist on sftp server")
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, ""
