def process(self, task):
detector = dtm.load_detector(task["detection_module"], task["session"])
results, report = detector.get_results()
# Store raw report to storage
Storage().store(task["scan_uuid"].hex, report)
# Change keys for conforming to result table schema
for result in results:
result["scan_id"] = task["scan_id"]
# Delete and insert scan results for keeping only the latest scan result
with db.database.atomic():
ResultTable.delete().where(
ResultTable.scan_id == task["scan_id"]).execute()
ResultTable.insert_many(results).execute()
ScanTable.update({
"ended_at": self.now
}).where(ScanTable.task_uuid == task["uuid"]).execute()
task["results"] = results
im.send(NotificationType.RESULT, task)
# Destroy the task without error
self.finish(task)
def post(self, audit_uuid):
"""
Create new scan into the specified audit
"""
params = Parser.ScanPostRequest.parse_args()
try:
detector = dtm.load_detector(params["detection_module"], None)
if detector.TARGET_TYPE == DetectionTarget.HOST.value:
validate_host(params["target"])
elif detector.TARGET_TYPE == DetectionTarget.URL.value:
params["target"] = get_safe_url(params["target"])
else:
abort(400, "Specified detector has invalid target type")
except Exception as e:
abort(400, str(e))
# Scan UUID consists of upper 96 bits audit UUID (=A) and lower 32 bits random number (=B),
# i.e., 'AAAAAAAA-AAAA-AAAA-AAAA-AAAABBBBBBBB'.
params["uuid"] = uuid.UUID(audit_uuid[0:24] + secrets.token_hex(4))
params["created_by"] = g.identity["name"]
params["updated_by"] = g.identity["name"]
audit, _ = get_audit_by_uuid(audit_uuid)
params["audit_id"] = audit["id"]
current_scan_count = ScanTable.select().where(
ScanTable.audit_id == params["audit_id"]).count()
if current_scan_count >= app.config["MAX_SCAN_COUNT_IN_EACH_AUDIT"]:
abort(400, "Max scan count exceeded")
ScanTable(**params).save()
return get_scan_by_uuid(params["uuid"])[0]
def process(self, task):
# Cancel if detector stays pending for a long time
created_at = task["created_at"].replace(tzinfo=pytz.utc)
if self.now > (created_at + timedelta(hours=app.config["SCAN_MAX_PENDING_DURATION_IN_HOUR"])):
raise Exception("Cancelled: detector stayed pending for a long time")
# Check if detector is ready to scan
detector = dtm.load_detector(task["detection_module"], task["session"])
if detector.is_ready():
# Enqueue the task to running queue
RunningTaskHandler().add(task)
def process(self, task):
# Cancel if detector stays running for a long time
started_at = task["started_at"].replace(tzinfo=pytz.utc)
if self.now > (started_at + timedelta(
hours=app.config["SCAN_MAX_RUNNING_DURATION_IN_HOUR"])):
raise Exception("Cancelled: detector was running for a long time")
# Check if detector is still running
detector = dtm.load_detector(task["detection_module"], task["session"])
if not detector.is_running():
# Enqueue the task to stopped queue
StoppedTaskHandler().add(task)
def add(self, task):
app.logger.info("Try to enqueue into {}: task={}".format(
self.progress, task))
detector = dtm.load_detector(task["detection_module"], task["session"])
session = detector.run(task["target"], task["detection_mode"])
task["session"] = session
task["progress"] = TaskProgress.RUNNING.name
task["started_at"] = self.now
with db.database.atomic():
# Update task progress
TaskTable.update(task).where(
TaskTable.uuid == task["uuid"]).execute()
# Reflect started time to corresponding scan entry
ScanTable.update({
"started_at": self.now
}).where(ScanTable.task_uuid == task["uuid"]).execute()
im.send(NotificationType.START, task)
app.logger.info("Enqueued into {} successfully: task={}".format(
self.progress, task["uuid"]))
return
def finish(self, task, error_reason=""):
app.logger.info("Try to delete: task={}, error_reason={}".format(
task, error_reason))
if task["session"] is not None:
detector = dtm.load_detector(task["detection_module"],
task["session"])
detector.delete()
with db.database.atomic():
TaskTable.delete().where(TaskTable.uuid == task["uuid"]).execute()
# Here we use task uuid for finding corresponding scan entry
# because task uuid of scan entry is changed by user's re-schedule
ScanTable.update({
"scheduled_at": None,
"task_uuid": None,
"error_reason": error_reason
}).where(ScanTable.task_uuid == task["uuid"]).execute()
if error_reason != "":
im.send(NotificationType.ERROR, task)
app.logger.info(
"Deleted successfully: task={}, error_reason={}".format(
task["uuid"], error_reason))
def add(self, scan):
app.logger.info("Try to enqueue into {}: scan={}".format(self.progress, scan))
detector = dtm.load_detector(scan["detection_module"], None)
if detector.TARGET_TYPE == DetectionTarget.HOST.value:
validate_host(scan["target"])
elif detector.TARGET_TYPE == DetectionTarget.URL.value:
scan["target"] = get_safe_url(scan["target"])
# Avoid concurrent scanning for the same target
task_query = TaskTable.select().where(TaskTable.target == scan["target"])
if task_query.count() > 0:
app.logger.info(
"Abandoned to enqueue scan={} because another scan for '{}' is still running".format(
scan["id"], scan["target"]
)
)
return None
session = detector.create()
task = {
"audit_id": scan["audit_id"],
"scan_id": scan["id"],
"scan_uuid": scan["uuid"],
"target": scan["target"],
"scheduled_at": scan["scheduled_at"],
"max_duration": scan["max_duration"],
"detection_module": scan["detection_module"],
"detection_mode": scan["detection_mode"],
"session": session,
"progress": TaskProgress.PENDING.name,
}
task = TaskTable(**task)
task.save()
app.logger.info("Enqueued into {} successfully: scan={}".format(self.progress, scan["id"]))
return task