def delete(self, scan_uuid): """ Delete the specified scan """ ScanTable.delete().where(ScanTable.uuid == scan_uuid).execute() return {}
def post(self, scan_uuid): """ Schedule the specified scan """ scan, _ = get_scan_by_uuid(scan_uuid) if scan["scheduled_at"] is not None: abort(400, "Scan is already scheduled") params = Parser.ScanSchedulePostRequest.parse_args() validate_schedule(params["scheduled_at"]) if len(params["rrule"]) > 0: scheduled_at = params["scheduled_at"].replace(tzinfo=pytz.utc) weekday = RRULE_WEEKDAY_LIST[scheduled_at.weekday()] hour = scheduled_at.time().hour params[ "rrule"] = "RRULE:FREQ=WEEKLY;BYDAY={};BYHOUR={};BYMINUTE=0;BYSECOND=0".format( weekday, hour) params["task_uuid"] = None params["started_at"] = None params["ended_at"] = None params["updated_by"] = g.identity["name"] ScanTable.update(params).where(ScanTable.uuid == scan_uuid).execute() return get_scan_by_uuid(scan_uuid)[0]
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): 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 set_next_scan(self): # Get scan entries that scheduled time has elapsed but still not be in any tasks scan_query = ( ScanTable().select().where((ScanTable.scheduled_at < fn.now()) & (ScanTable.task_uuid.is_null())) ) for scan in scan_query.dicts(): try: app.logger.info("Try to set: scan={}".format(scan)) # Cancel scan if scan period has already elapsed scheduled_at = scan["scheduled_at"].replace(tzinfo=pytz.utc) if self.now > (scheduled_at + timedelta(hours=scan["max_duration"])): raise Exception("Cancelled: scheduled period has elapsed") with db.database.atomic(): # Enqueue the task to pending queue task = PendingTaskHandler().add(scan) if task: # ToDo: Consider race condition between scan reschedule API and this thread context ScanTable.update({"task_uuid": task.uuid}).where(ScanTable.id == scan["id"]).execute() app.logger.info("Set scan successfully: scan={}".format(scan["id"])) except Exception as error: app.logger.warn("ERROR: scan={}, error={}".format(scan["id"], error)) self.reset_scan_schedule(scan, error)
def set_next_periodic_scan_schedule(self): scan_query = ScanTable().select().where((ScanTable.scheduled_at.is_null() & (ScanTable.rrule != ""))) for scan in scan_query.dicts(): app.logger.info("Try to schedule next: scan={}".format(scan)) try: next_schedule = rrulestr(scan["rrule"])[0].replace(tzinfo=pytz.utc) ScanTable.update({"scheduled_at": next_schedule}).where(ScanTable.id == scan["id"]).execute() app.logger.info("Scheduled next successfully: scan={}".format(scan["id"])) except Exception as error: app.logger.error("ERROR: scan={}, error={}".format(scan["id"], error))
def send(self, notification_type, task): integrations = IntegrationTable.select().where( IntegrationTable.audit_id == task["audit_id"]) if len(integrations.dicts()) > 0: scan = ScanTable.select().where( ScanTable.id == task["scan_id"]).dicts()[0] for integration in integrations.dicts(): self.integrators[integration["service"]]().send( notification_type, scan, task, integration)
def delete(self, scan_uuid): """ Cancel the specified scan schedule """ scan, _ = get_scan_by_uuid(scan_uuid) if scan["scheduled_at"] is None: abort(400, "Scan is not scheduled") default_values = { "scheduled_at": None, "max_duration": 0, "rrule": "", "started_at": None, "ended_at": None, "task_uuid": None, "updated_by": g.identity["name"], } ScanTable.update(default_values).where( ScanTable.uuid == scan_uuid).execute() return get_scan_by_uuid(scan_uuid)[0]
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 reset_scan_schedule(self, scan, error_reason=""): ScanTable.update({"scheduled_at": None, "error_reason": error_reason, "task_uuid": None}).where( ScanTable.id == scan["id"] ).execute()
def get_scan_by_uuid(scan_uuid): try: query = ScanTable.select().where(ScanTable.uuid == scan_uuid) return query.dicts()[0], query except: abort(404)