Merge pull request #5 from GitHubNull/support_progress_num_arg

support process num arg for sqlmapapi.

Author: GitHubNull

lib/utils/api.py

@@ -355,8 +355,11 @@ def is_admin(token):
     return DataStore.admin_token == token
 
 
-def perform_task():
+def perform_task(num_cores=None):
     # 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:
@@ -379,9 +382,9 @@ def perform_task():
                 else:
                     running_task_count += 1
 
-        if running_task_count < MAX_TASKS_NUMBER:
+        if running_task_count < local_max_tasks_number:
             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 datetime.datetime.now() >= task.start_datetime:
                             running_task_count += 1
@@ -398,22 +401,22 @@ def perform_task():
                         task.status = TaskStatus.Running
 
 
-def run_task(interval):
+def run_task(interval, num_cores):
     logger.debug("run_task...")
     try:
         while True:
             # 执行定时任务
-            perform_task()
+            perform_task(num_cores)
             # 等待一定时间
             time.sleep(interval)
     except KeyboardInterrupt:
         print("定时任务已停止")
 
 
-def schedule_task(interval):
+def schedule_task(interval, num_cores):
     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)
     # 启动线程
@@ -1190,7 +1193,7 @@ def version(token=None):
     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
     """
@@ -1218,7 +1221,7 @@ def server(host=RESTAPI_DEFAULT_ADDRESS, port=RESTAPI_DEFAULT_PORT, adapter=REST
     DataStore.current_db.init()
 
     # 开启定时任务
-    schedule_task(1)
+    schedule_task(1, num_cores)
 
     # Run RESTful API
     try:

sqlmapapi.py

@@ -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("--username", help="Basic authentication username (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()
 
     # Start the client or the 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:
         client(args.host, args.port, username=args.username, password=args.password)
     else: