sqlmapproject/sqlmap
1
1
Fork 0
mirror of https://github.com/sqlmapproject/sqlmap.git synced 2025-07-26 07:59:52 +03:00
Code Issues Packages Projects Releases Wiki Activity

support process num arg for sqlmapapi.

Browse Source
This commit is contained in:
Name 2024-01-06 22:53:59 +08:00
parent c813622b49
commit a0a90ebc05
2 changed files with 16 additions and 10 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

21
lib/utils/api.py
View File

@ -355,8 +355,11 @@ def is_admin(token):
return DataStore.admin_token == token return DataStore.admin_token == token
def perform_task(): def perform_task(num_cores=None):
# logger.debug('perform_task...') # logger.debug('perform_task...')
global MAX_TASKS_NUMBER
local_max_tasks_number = MAX_TASKS_NUMBER if num_cores is None or num_cores == 0 else num_cores
# 计算在扫描的任务的数量 # 计算在扫描的任务的数量
with DataStore.tasks_lock: with DataStore.tasks_lock:
@ -379,9 +382,9 @@ def perform_task():
else: else:
running_task_count += 1 running_task_count += 1
if running_task_count < MAX_TASKS_NUMBER: if running_task_count < local_max_tasks_number:
for task in runnable_list: for task in runnable_list:
if running_task_count < MAX_TASKS_NUMBER: if running_task_count < local_max_tasks_number:
if task.start_datetime is not None: if task.start_datetime is not None:
if datetime.datetime.now() >= task.start_datetime: if datetime.datetime.now() >= task.start_datetime:
running_task_count += 1 running_task_count += 1
@ -398,22 +401,22 @@ def perform_task():
task.status = TaskStatus.Running task.status = TaskStatus.Running
def run_task(interval): def run_task(interval, num_cores):
logger.debug("run_task...") logger.debug("run_task...")
try: try:
while True: while True:
# 执行定时任务 # 执行定时任务
perform_task() perform_task(num_cores)
# 等待一定时间 # 等待一定时间
time.sleep(interval) time.sleep(interval)
except KeyboardInterrupt: except KeyboardInterrupt:
print("定时任务已停止") print("定时任务已停止")
def schedule_task(interval): def schedule_task(interval, num_cores):
logger.debug("schedule_task...") logger.debug("schedule_task...")
# 创建后台线程 # 创建后台线程
thread = threading.Thread(target=run_task, args=(interval,)) thread = threading.Thread(target=run_task, args=(interval, num_cores, ))
# 设置线程为守护线程 # 设置线程为守护线程
thread.setDaemon(True) thread.setDaemon(True)
# 启动线程 # 启动线程
@ -1190,7 +1193,7 @@ def version(token=None):
return jsonize({"success": True, "version": VERSION_STRING.split('/')[-1]}) return jsonize({"success": True, "version": VERSION_STRING.split('/')[-1]})
def server(host=RESTAPI_DEFAULT_ADDRESS, port=RESTAPI_DEFAULT_PORT, adapter=RESTAPI_DEFAULT_ADAPTER, username=None, password=None): def server(host=RESTAPI_DEFAULT_ADDRESS, port=RESTAPI_DEFAULT_PORT, adapter=RESTAPI_DEFAULT_ADAPTER, username=None, password=None, num_cores=None):
""" """
REST-JSON API server REST-JSON API server
""" """
@ -1218,7 +1221,7 @@ def server(host=RESTAPI_DEFAULT_ADDRESS, port=RESTAPI_DEFAULT_PORT, adapter=REST
DataStore.current_db.init() DataStore.current_db.init()
# 开启定时任务 # 开启定时任务
schedule_task(1) schedule_task(1, num_cores)
# Run RESTful API # Run RESTful API
try: try:

5
sqlmapapi.py
View File

@ -60,11 +60,14 @@ def main():
apiparser.add_option("--adapter", help="Server (bottle) adapter to use (default \"%s\")" % RESTAPI_DEFAULT_ADAPTER, default=RESTAPI_DEFAULT_ADAPTER, action="store") apiparser.add_option("--adapter", help="Server (bottle) adapter to use (default \"%s\")" % RESTAPI_DEFAULT_ADAPTER, default=RESTAPI_DEFAULT_ADAPTER, action="store")
apiparser.add_option("--username", help="Basic authentication username (optional)", action="store") apiparser.add_option("--username", help="Basic authentication username (optional)", action="store")
apiparser.add_option("--password", help="Basic authentication password (optional)", action="store") apiparser.add_option("--password", help="Basic authentication password (optional)", action="store")
# Add an option to specify the number of CPU cores
apiparser.add_option("-n", "--num-cores", help="Number of CPU cores to use", type="int", action="store", default=0)
(args, _) = apiparser.parse_args() (args, _) = apiparser.parse_args()
# Start the client or the server # Start the client or the server
if args.server: if args.server:
server(args.host, args.port, adapter=args.adapter, username=args.username, password=args.password) server(args.host, args.port, adapter=args.adapter, username=args.username, password=args.password, num_cores=args.num_cores)
elif args.client: elif args.client:
client(args.host, args.port, username=args.username, password=args.password) client(args.host, args.port, username=args.username, password=args.password)
else: else: