update: enhance api logical

Basis/Api.py

@@ -8,6 +8,7 @@ from Checker import formatCheck
 from Basis.Logger import logging
 from Basis.Manager import Manager
 from flask import Flask, Response, request
+from Basis.Exception import managerException
 from Basis.Constant import ApiPort, ApiPath, ApiToken, Version
 
 webApi = Flask(__name__)  # init flask server
@@ -68,12 +69,17 @@ def createTask() -> Response:
     if not tokenCheck():  # token check
         return genError('Invalid token')
 
-    # TODO: format check and proxy list
-    checkList = formatCheck(request.json.get('check'))
+    try:
+        # TODO: format check and proxy list
+        checkList = formatCheck(request.json.get('check'))
+    except:
+        return genError('Some error in check options')
     proxyList = []
     for proxy in request.json.get('proxy'):
-        proxyList.append(formatProxy(proxy))
-    logging.critical(proxyList)
+        try:
+            proxyList.append(formatProxy(proxy))
+        except Exception as exp:
+            return genError('Proxy error in %s -> %s' % (proxy, exp))
 
     logging.debug('API create task -> check = %s | proxy = %s' % (checkList, proxyList))
     tasks = []
@@ -105,6 +111,22 @@ def getTaskInfo(taskId: str) -> Response:
     })
 
 
+@webApi.route(os.path.join(ApiPath, 'task/<taskId>'), methods = ['DELETE'])
+def deleteTask(taskId: str) -> Response:
+    if not tokenCheck():  # token check
+        return genError('Invalid token')
+    logging.debug('API get task -> %s' % taskId)
+    if not Manager.isUnion(taskId):
+        return genError('Task not found')
+    try:
+        Manager.delUnion(taskId)
+        return jsonResponse({
+            'success': True
+        })
+    except managerException as exp:
+        return genError(str(exp))
+
+
 @webApi.route(os.path.join(ApiPath, 'version'), methods = ['GET'])
 def getVersion() -> Response:
     logging.debug('API get version -> %s' + Version)

Basis/Check.py

@@ -6,6 +6,8 @@ import time
 from Checker import Checker
 from Basis.Logger import logging
 from Builder import Builder, clientEntry
+from Basis.Exception import checkException
+from Basis.Functions import checkPortStatus
 
 
 def buildClient(taskId: str, taskInfo: dict) -> Builder:
@@ -18,24 +20,27 @@ def buildClient(taskId: str, taskInfo: dict) -> Builder:
         )
     except Exception as reason:
         logging.error('[%s] Client build error -> %s' % (taskId, reason))
-        raise RuntimeError('Client build error')
+        raise checkException('Client build error')
 
 
-def waitClient(taskId: str, client: Builder):
-    # TODO: wait port occupied (client.socksPort)
-    time.sleep(1)  # TODO: simple delay for now
+def waitClient(taskId: str, client: Builder, times: int = 150, delay: int = 100):  # wait until client port occupied
+    for i in range(times):
+        if not checkPortStatus(client.socksPort):  # port occupied
+            break
+        time.sleep(delay / 1000)  # wait in default: 100ms * 150 => 15s
+    time.sleep(1)  # wait a short time before check process
     if not client.status():  # client unexpected exit
         logging.warning('[%s] Client unexpected exit' % taskId)
         client.destroy()  # remove file and kill sub process
         logging.debug('[%s] Client output\n%s', (taskId, client.output))
-        raise RuntimeError('Client unexpected exit')
+        raise checkException('Client unexpected exit')
 
 
 def Check(taskId: str, taskInfo: dict) -> dict:
     logging.info('[%s] Start checking process -> %s' % (taskId, taskInfo))
     if taskInfo['type'] not in clientEntry:  # unknown proxy type
         logging.error('[%s] Unknown proxy type %s' % (taskId, taskInfo['type']))
-        raise RuntimeError('Unknown proxy type')
+        raise checkException('Unknown proxy type')
     client = buildClient(taskId, taskInfo)  # build proxy client
     logging.info('[%s] Client loaded successfully' % taskId)
     waitClient(taskId, client)  # wait for the client to start
@@ -49,5 +54,6 @@ def Check(taskId: str, taskInfo: dict) -> dict:
     taskInfo.pop('check')  # remove check items
     return {
         **taskInfo,
+        'success': True,
         'result': checkResult,  # add check result
     }

Basis/Constant.py

@@ -20,10 +20,15 @@ TestHost = 'proxyc.net'
 TestSite = 'www.bing.com'
 PathEnv = '/bin:/sbin:/usr/bin:/usr/sbin:/usr/local/bin'
 
+
 # Load Env Options
-yamlFile = os.path.join(os.path.dirname(os.path.realpath(__file__)), '../env.yaml')
-yamlContent = open(yamlFile, 'r', encoding = 'utf-8').read()
-envOptions = yaml.load(yamlContent, Loader = yaml.FullLoader)
+envOptions = {}
+try:
+    yamlFile = os.path.join(os.path.dirname(os.path.realpath(__file__)), '../env.yaml')
+    yamlContent = open(yamlFile, 'r', encoding = 'utf-8').read()
+    envOptions = yaml.load(yamlContent, Loader = yaml.FullLoader)
+except:  # something error in env.yaml
+    pass
 if 'version' in envOptions:
     Version = envOptions['version']
 if 'loglevel' in envOptions:
@@ -40,12 +45,14 @@ if 'api' in envOptions:
     if 'token' in envOptions['api']:
         ApiToken = envOptions['api']['token']
 
+
 # WorkDir Create
 try:
     os.makedirs(WorkDir)  # just like `mkdir -p ...`
 except:
     pass  # folder exist or target is another thing
 
+
 # Shadowsocks Info
 mbedtlsMethods = [
     'aes-128-cfb128',
@@ -120,6 +127,7 @@ ssAllMethods = set()
 [ssAllMethods.update(ssMethods[x]) for x in ssMethods]
 ssAllMethods = sorted(list(ssAllMethods))  # methods of Shadowsocks
 
+
 # Plugin Info
 Plugins = {
     'simple-obfs': ['obfs-local', 'obfs-server'],
@@ -140,6 +148,7 @@ Plugins = {x: [Plugins[x][0], Plugins[x][1 if len(Plugins[x]) == 2 else 0]] for
 Plugins = {x: {'client': Plugins[x][0], 'server': Plugins[x][1]} for x in Plugins}  # format plugins info
 pluginClients = [Plugins[x]['client'] for x in Plugins]  # plugin client list -> obfs-local / simple-tls / ...
 
+
 # ShadowsocksR Info
 ssrMethods = [  # methods of ShadowsocksR
     'aes-128-ctr', 'aes-192-ctr', 'aes-256-ctr',
@@ -168,19 +177,20 @@ ssrObfuscations = [ # obfuscations of ShadowsocksR (obfs)
     'tls_simple', 'tls1.2_ticket_auth', 'tls1.2_ticket_fastauth',
 ]
 
-# VMess Info
+
+# V2ray / Xray Info
 vmessMethods = ['aes-128-gcm', 'chacha20-poly1305', 'auto', 'none', 'zero']
 
-# XTLS Info
+quicMethods = ['none', 'aes-128-gcm', 'chacha20-poly1305']
+udpObfuscations = ['none', 'srtp', 'utp', 'wechat-video', 'dtls', 'wireguard']
+
 xtlsFlows = ['xtls-origin', 'xtls-direct', 'xtls-splice']
 xtlsFlows = {x: x.replace('-', '-rprx-') for x in xtlsFlows}
 
-# v2ray / Xray Info
-quicMethods = ['none', 'aes-128-gcm', 'chacha20-poly1305']
-udpObfuscations = ['none', 'srtp', 'utp', 'wechat-video', 'dtls', 'wireguard']
 
 # Trojan-Go Info
 trojanGoMethods = ['aes-128-gcm', 'aes-256-gcm', 'chacha20-ietf-poly1305']
 
+
 # Hysteria Info
 hysteriaProtocols = ['udp', 'wechat-video', 'faketcp']

Basis/Manager.py

@@ -39,15 +39,26 @@ class Task(object):
             logging.info('Manager add task [%s] -> %s' % (taskId, task))
         self.__tasks.update(tasks)  # load into task list
         self.__unions[unionId] = {
+            'finish': False,
             'items': taskIds  # record task items
         }
         logging.info('Manager add union [%s] -> %s' % (unionId, taskIds))
         return unionId
 
+    def delUnion(self, unionId) -> None:  # remove union
+        if unionId not in self.__unions:
+            logging.error('Manager union [%s] not found' % unionId)
+            raise managerException('Union id not found')
+        if not self.__unions[unionId]['finish']:  # some tasks are still running
+            raise managerException('Couldn\'t remove working union')
+        self.__unions.pop(unionId)
+
     def getUnion(self, unionId: str) -> dict:  # get union status (remove tasks when all completed)
         if unionId not in self.__unions:
             logging.error('Manager union [%s] not found' % unionId)
             raise managerException('Union id not found')
+        if self.__unions[unionId]['finish']:  # all tasks are finished
+            return self.__unions[unionId]
         tasks = self.__unions[unionId]['items']
         finishNum = 0
         for taskId in tasks:
@@ -60,17 +71,15 @@ class Task(object):
                 'finish': False,
                 'percent': round(finishNum / len(tasks), 2)
             }
-        self.__unions.pop(unionId)  # remove from union list
-        unionResult = []  # temporary storage
+        self.__unions[unionId]['result'] = []
         for taskId in tasks:
             task = self.__tasks[taskId]
             self.__tasks.pop(taskId)  # remove from task list
-            unionResult.append(task['data'])
-        logging.info('Manager release union [%s] -> %s' % (unionId, unionResult))
-        return {
-            'finish': True,
-            'result': unionResult
-        }
+            self.__unions[unionId]['result'].append(task['data'])
+        self.__unions[unionId]['finish'] = True
+        self.__unions[unionId].pop('items')
+        logging.info('Manager release union [%s] -> %s' % (unionId, self.__unions[unionId]['result']))
+        return self.__unions[unionId]
 
     def popTask(self) -> tuple[str or None, any]:  # fetch a loaded task
         for taskId, task in self.__tasks.items():

main.py

@@ -8,6 +8,7 @@ import _thread
 import argparse
 import compileall
 from Basis import Constant
+from Basis.Exception import checkException
 
 
 def mainArgParse(rawArgs: list) -> argparse.Namespace:
@@ -74,9 +75,25 @@ def pythonCompile(dirRange: str = '/') -> None:  # python optimize compile
 
 
 def runCheck(taskId: str, taskInfo: dict) -> None:
-    checkResult = Check(taskId, taskInfo)  # check by task info
-    logging.warning('[%s] Task finish' % taskId)
-    Manager.finishTask(taskId, checkResult)  # commit check result
+    success = True
+    checkResult = {}
+    try:
+        checkResult = Check(taskId, taskInfo)  # check by task info
+        logging.warning('[%s] Task finish' % taskId)
+    except checkException as exp:
+        success = False
+        logging.error('[%s] Task error -> %s' % (taskId, exp))
+    except:
+        success = False
+        logging.error('[%s] Task error -> Unknown error' % taskId)
+    finally:
+        if not success:  # got some error in check process
+            taskInfo.pop('check')
+            checkResult = {
+                **taskInfo,
+                'success': False,
+            }
+        Manager.finishTask(taskId, checkResult)  # commit check result
 
 
 def loop(threadNum: int = 16) -> None: