add start datetime and update start datetime feature

This commit is contained in:
Name 2023-12-24 15:08:08 +08:00
parent 2ae56158e4
commit 452780e792

View File

@ -9,6 +9,7 @@ See the file 'LICENSE' for copying permission
from __future__ import print_function from __future__ import print_function
import contextlib import contextlib
import datetime
import logging import logging
import os import os
import re import re
@ -60,6 +61,7 @@ from lib.core.subprocessng import Popen
from lib.parse.cmdline import cmdLineParser from lib.parse.cmdline import cmdLineParser
from thirdparty.bottle.bottle import error as return_error from thirdparty.bottle.bottle import error as return_error
from thirdparty.bottle.bottle import get from thirdparty.bottle.bottle import get
from thirdparty.bottle.bottle import route
from thirdparty.bottle.bottle import hook from thirdparty.bottle.bottle import hook
from thirdparty.bottle.bottle import post from thirdparty.bottle.bottle import post
from thirdparty.bottle.bottle import request from thirdparty.bottle.bottle import request
@ -78,6 +80,7 @@ from lib.utils.task_status_enum import TaskStatus
# Global data storage # Global data storage
MAX_TASKS_NUMBER = multiprocessing.cpu_count() - 1 MAX_TASKS_NUMBER = multiprocessing.cpu_count() - 1
ROOT_DIRECTORY = os.getcwd() ROOT_DIRECTORY = os.getcwd()
datetime_format = "%Y-%m-%d %H:%M:%S"
class DataStore(object): class DataStore(object):
admin_token = "" admin_token = ""
@ -169,6 +172,7 @@ class Task(object):
self.options = None self.options = None
self.status = TaskStatus.New self.status = TaskStatus.New
self._original_options = None self._original_options = None
self.start_datetime = None
self.initialize_options(taskid) self.initialize_options(taskid)
def initialize_options(self, taskid): def initialize_options(self, taskid):
@ -378,10 +382,20 @@ def perform_task():
if running_task_count < MAX_TASKS_NUMBER: if running_task_count < MAX_TASKS_NUMBER:
for task in runnable_list: for task in runnable_list:
if running_task_count < MAX_TASKS_NUMBER: if running_task_count < MAX_TASKS_NUMBER:
running_task_count += 1 if task.start_datetime is not None:
logger.info("run task %s" % task.options.taskid) if datetime.datetime.now() >= task.start_datetime:
task.engine_start() running_task_count += 1
task.status = TaskStatus.Running logger.info("run task %s" % task.options.taskid)
task.engine_start()
task.status = TaskStatus.Running
else:
continue
else:
running_task_count += 1
logger.info("run task %s" % task.options.taskid)
task.start_datetime = datetime.datetime.now()
task.engine_start()
task.status = TaskStatus.Running
def run_task(interval): def run_task(interval):
@ -441,6 +455,8 @@ def security_headers(json_header=True):
response.headers["X-XSS-Protection"] = "1; mode=block" response.headers["X-XSS-Protection"] = "1; mode=block"
response.headers["Pragma"] = "no-cache" response.headers["Pragma"] = "no-cache"
response['Access-Control-Allow-Origin'] = 'http://localhost:5173' response['Access-Control-Allow-Origin'] = 'http://localhost:5173'
response.headers['Access-Control-Allow-Methods'] = 'PUT, GET, POST, DELETE, OPTIONS'
response.headers['Access-Control-Allow-Headers'] = 'Origin, Accept, Content-Type, X-Requested-With, X-CSRF-Token'
response.headers["Cache-Control"] = "no-cache" response.headers["Cache-Control"] = "no-cache"
response.headers["Expires"] = "0" response.headers["Expires"] = "0"
@ -449,6 +465,11 @@ def security_headers(json_header=True):
# else: # else:
# response.content_type = "text/html; charset=utf-8" # response.content_type = "text/html; charset=utf-8"
# 处理 OPTIONS 请求
@route('/<path:path>', method=['OPTIONS'])
def options_handler(path):
return
############################## ##############################
# HTTP Status Code functions # # HTTP Status Code functions #
############################## ##############################
@ -648,6 +669,7 @@ def task_ls(token=None):
resul_task_item = { resul_task_item = {
"index": index, "index": index,
"start_datetime": None if task.start_datetime is None else task.start_datetime.strftime("%Y-%m-%d %H:%M:%S"),
"task_id": taskid, "task_id": taskid,
"errors": errors_count, "errors": errors_count,
"logs": logs_count, "logs": logs_count,
@ -792,6 +814,114 @@ def scan_start(taskid):
logger.debug("Add [%s] to scan list" % taskid) logger.debug("Add [%s] to scan list" % taskid)
return jsonize({"success": True, "engineid": 0}) return jsonize({"success": True, "engineid": 0})
@post('/scan/start_at_datetime/<taskid>')
def scan_start_at_datetime(taskid):
"""
Start a scan at a specific datetime
"""
with DataStore.tasks_lock:
if taskid not in DataStore.tasks:
logger.warning("[%s] Invalid task ID provided to scan_start()" % taskid)
return jsonize({"success": False, "message": "Invalid task ID"})
if request.json is None:
return jsonize({"success": False, "message": "Invalid request"})
params = request.params
if 'start_datetime' not in params:
return jsonize({"success": False, "message": "Invalid start_datetime"})
start_datetime = params['start_datetime']
if not isinstance(start_datetime, str):
return jsonize({"success": False, "message": "Invalid start_datetime"})
for key in request.json:
if key in RESTAPI_UNSUPPORTED_OPTIONS:
logger.warning(
"[%s] Unsupported option '%s' provided to scan_start()" % (taskid, key))
return jsonize({"success": False, "message": "Unsupported option '%s'" % key})
with DataStore.tasks_lock:
if DataStore.tasks[taskid].status == TaskStatus.Blocked:
DataStore.tasks[taskid].status = TaskStatus.Runnable
logger.debug("(%s) Unblocked" % taskid)
return jsonize({"success": True, "engineid": 0})
for option, value in request.json.items():
DataStore.tasks[taskid].set_option(option, value)
# Launch sqlmap engine in a separate process
DataStore.tasks[taskid].status = TaskStatus.Runnable
DataStore.tasks[taskid].start_datetime = datetime.datetime.strptime(start_datetime, datetime_format)
logger.debug("Add (%s) to scan list" % taskid)
return jsonize({"success": True, "engineid": 0})
@post('/scan/update_start_datetime/<taskid>')
def scan_update_start_datetime(taskid):
"""
Update the start datetime of a scan
"""
logger.debug("[%s] Updating start datetime" % taskid)
with DataStore.tasks_lock:
if taskid not in DataStore.tasks:
logger.warning("[%s] Invalid task ID provided to scan_update_start_datetime()" % taskid)
return jsonize({"success": False, "message": "Invalid task ID"})
start_datetime = request.json.get("start_datetime", None)
if start_datetime is None:
logger.warning("[%s] No start_datetime provided to scan_update_start_datetime()" % taskid)
return jsonize({"success": False, "message": "Invalid start datetime"})
now = datetime.datetime.now()
time_five_seconds_later = now + datetime.timedelta(seconds=5)
start_datetime = datetime.datetime.strptime(start_datetime, datetime_format)
if DataStore.tasks[taskid].start_datetime is None:
if start_datetime > time_five_seconds_later:
DataStore.tasks[taskid].start_datetime = start_datetime
return jsonize({"success": True, "message": "update start da"})
else:
return jsonize({"success": False, "message": "start datetime is too early"})
else:
if DataStore.tasks[taskid].status in [TaskStatus.New, TaskStatus.Runnable]:
if start_datetime > now:
DataStore.tasks[taskid].start_datetime = start_datetime
return jsonize({"success": True, "message": "update start datetime success"})
else:
return jsonize({"success": False, "message": "start datetime must be greater than now"})
elif DataStore.tasks[taskid].status == TaskStatus.Running:
# 检查你的datetime对象是否大于现在时间5秒
if start_datetime > time_five_seconds_later:
DataStore.tasks[taskid].engine_stop()
DataStore.tasks[taskid].start_datetime = start_datetime
return jsonize({"success": True, "message": "Update start datetime success"})
else:
return jsonize({"success": False, "message": "Invalid start datetime"})
elif DataStore.tasks[taskid].status == TaskStatus.Terminated:
if start_datetime > time_five_seconds_later:
DataStore.tasks[taskid].start_datetime = start_datetime
DataStore.tasks[taskid].status = TaskStatus.Runnable
return jsonize({"success": True, "message": "Task resumed"})
else:
return jsonize({"success": False, "message": "Invalid start datetime"})
elif DataStore.tasks[taskid].status == TaskStatus.Blocked:
if start_datetime > time_five_seconds_later:
DataStore.tasks[taskid].start_datetime = start_datetime
DataStore.tasks[taskid].status = TaskStatus.Runnable
return jsonize({"success": True, "message": "Task resumed"})
else:
return jsonize({"success": False, "message": "Invalid start datetime"})
else:
return jsonize({"success": False, "message": "Invalid task status"})
@get('/scan/startBlocked/<taskid>') @get('/scan/startBlocked/<taskid>')
def scan_startBlocked(taskid): def scan_startBlocked(taskid):
""" """