def stopTask(pentest, tool_iid, body):
mongoInstance = MongoCalendar.getInstance()
mongoInstance.connectToDb(pentest)
stopableTool = ServerTool.fetchObject(pentest, {"_id": ObjectId(tool_iid)})
print("Trying to stop task " + str(stopableTool))
if stopableTool is None:
return "Tool not found", 404
workers = mongoInstance.getWorkers({})
workerNames = [worker["name"] for worker in workers]
forceReset = body["forceReset"]
saveScannerip = stopableTool.scanner_ip
if forceReset:
stopableTool.markAsNotDone()
update(pentest, tool_iid, ToolController(stopableTool).getData())
if saveScannerip == "":
return "Empty worker field", 400
if saveScannerip == "localhost":
return "Tools running in localhost cannot be stopped through API", 405
if saveScannerip not in workerNames:
return "The worker running this tool is not running anymore", 404
from pollenisator.api import socketio, sockets
socketio.emit('stopCommand', {
'pentest': pentest,
"tool_iid": str(tool_iid)
},
room=sockets[saveScannerip])
if not forceReset:
stopableTool.markAsNotDone()
update(pentest, tool_iid, ToolController(stopableTool).getData())
return "Success", 200
def setInTime(self):
"""Set this tool as in time (matching any interval in wave)
Remove "OOT" from status
"""
if "OOT" in self.status:
self.status.remove("OOT")
update(self.pentest, self._id, ToolController(self).getData())
def setInScope(self):
"""Set this tool as out of scope (not matching any scope in wave)
Add "OOS" in status
"""
if "OOS" in self.status:
self.status.remove("OOS")
update(self.pentest, self._id, ToolController(self).getData())
def importResult(pentest, tool_iid, upfile, body):
#STORE FILE
plugin = body.get("plugin", "")
res, status, filepath = _upload(pentest, tool_iid, "result", upfile)
if status != 200:
return res, status
# Analyze
upfile.stream.seek(0)
mongoInstance = MongoCalendar.getInstance()
mongoInstance.connectToDb(pentest)
toolModel = ServerTool.fetchObject(pentest, {"_id": ObjectId(tool_iid)})
if toolModel is None:
return "Tool not found", 404
mod = toolModel.getPlugin(plugin)
if mod is not None:
try:
# Check return code by plugin (can be always true if the return code is inconsistent)
notes, tags, _, _ = mod.Parse(pentest, upfile)
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.notes = notes
toolModel.tags = tags
toolModel.markAsDone(filepath)
# And update the tool in database
update(pentest, tool_iid, ToolController(toolModel).getData())
# Upload file to SFTP
msg = "TASK SUCCESS : " + toolModel.name
except IOError as e:
toolModel.tags = ["todo"]
toolModel.notes = "Failed to read results file"
toolModel.markAsDone()
update(pentest, tool_iid, ToolController(toolModel).getData())
else:
msg = "TASK FAILED (no plugin found) : " + toolModel.name
toolModel.markAsNotDone()
update(pentest, tool_iid, ToolController(toolModel).getData())
raise Exception(msg)
return "Success"
def unregister(name):
worker = mongoInstance.getWorker(name)
if worker is not None:
running_tools = worker.get("running_tools", [])
for running_tool in running_tools:
tool_m = ServerTool.fetchObject(running_tool["pentest"], {"_id":ObjectId(running_tool["iid"])})
if "running" in tool_m.getStatus():
tool_m.markAsNotDone()
tool_update(running_tool["pentest"], running_tool["iid"], ToolController(tool_m).getData())
mongoInstance.deleteFromDb("pollenisator", "workers", {"name": name}, False, True)
return True
return "Worker not Found", 404
def removeWorkers():
workers = mongoInstance.getWorkers()
count = 0
for worker in workers:
running_tools = worker.get("running_tools", [])
for running_tool in running_tools:
tool_m = ServerTool.fetchObject(running_tool["pentest"], {"_id":ObjectId(running_tool["iid"])})
if "running" in tool_m.getStatus():
tool_m.markAsNotDone()
tool_update(running_tool["pentest"], running_tool["iid"], ToolController(tool_m).getData())
deleteWorker(worker["name"])
count += 1
return {"n":int(count)}
def launchTask(pentest, tool_iid, body, **kwargs):
worker_token = kwargs.get("worker_token") if kwargs.get(
"worker_token") else encode_token(kwargs.get("token_info"))
mongoInstance = MongoCalendar.getInstance()
mongoInstance.connectToDb(pentest)
launchableTool = ServerTool.fetchObject(pentest,
{"_id": ObjectId(tool_iid)})
if launchableTool is None:
return "Tool not found", 404
checks = body["checks"]
plugin = body["plugin"]
# Find a worker that can launch the tool without breaking limitations
workers = mongoInstance.getWorkers({"pentests": pentest})
choosenWorker = ""
for worker in workers:
workerName = worker["name"]
if hasRegistered(workerName, launchableTool):
if not checks:
choosenWorker = workerName
elif hasSpaceFor(workerName, launchableTool, pentest):
choosenWorker = workerName
break
if choosenWorker == "":
return "No worker available", 404
workerName = choosenWorker
launchableToolId = launchableTool.getId()
launchableTool.markAsRunning(workerName)
update(pentest, tool_iid, ToolController(launchableTool).getData())
# Mark the tool as running (scanner_ip is set and dated is set, datef is "None")
# Use socket sid as room so that only this worker will receive this task
from pollenisator.api import socketio, sockets
socketio.emit('executeCommand', {
'workerToken': worker_token,
"pentest": pentest,
"toolId": str(launchableToolId),
"parser": plugin
},
room=sockets[workerName])
return "Success ", 200
def setStatus(pentest, tool_iid, body):
newStatus = body["newStatus"]
arg = body.get("arg", "")
tool_o = ServerTool.fetchObject(pentest, {"_id": ObjectId(tool_iid)})
if tool_o is None:
return "Tool not found", 404
if "done" in newStatus:
if arg == "":
arg = None
tool_o.markAsDone(arg)
elif "running" in newStatus:
tool_o.markAsRunning(arg)
elif "not_done" in newStatus:
tool_o.markAsNotDone()
elif "ready" in newStatus:
tool_o.markAsNotDone()
elif "error" in newStatus:
tool_o.markAsError()
elif "timedout" in newStatus:
tool_o.markAsTimedout()
elif len(newStatus) == 0:
tool_o.markAsNotDone()
return update(pentest, tool_o.getId(), ToolController(tool_o).getData())
def addInDb(self):
return insert(self.pentest, ToolController(self).getData())