update: optimize code

Basis/Api.py

@@ -17,10 +17,10 @@ def jsonResponse(data: dict) -> Response:  # return json mime
     return Response(json.dumps(data), mimetype = 'application/json')
 
 
-def tokenError() -> Response:
+def genError(message: str) -> Response:
     return jsonResponse({
         'success': False,
-        'message': 'Invalid token'
+        'message': message,
     })
 
 
@@ -37,47 +37,33 @@ def tokenCheck() -> bool:
 @webApi.route('/task', methods = ['GET'])
 def getTaskList() -> Response:
     if not tokenCheck():  # token check
-        return tokenError()
-    try:
-        taskList = Manager.listUnion()
-        logging.debug('api get task list -> %s' % taskList)
-        return jsonResponse({
-            'success': True,
-            'task': taskList,
-        })
-    except:
-        return jsonResponse({
-            'success': False,
-            'message': 'Unknown error'
-        })
+        return genError('Invalid token')
+    taskList = Manager.listUnion()
+    logging.debug('API get task list -> %s' % taskList)
+    return jsonResponse({
+        'success': True,
+        'task': taskList,
+    })
 
 
 @webApi.route('/task', methods = ['POST'])
 def createTask() -> Response:
     if not tokenCheck():  # token check
-        return tokenError()
-    checkList = request.json.get('check')
-    proxyList = request.json.get('proxy')
-    if checkList is None or type(checkList) != list:
-        return jsonResponse({
-            'success': False,
-            'message': 'invalid check list',
-        })
-    if proxyList is None or type(proxyList) != list:
-        return jsonResponse({
-            'success': False,
-            'message': 'invalid proxy list',
-        })
-    logging.debug('api create task -> check: %s | proxy: %s' % (checkList, proxyList))
+        return genError('Invalid token')
 
     # TODO: format check and proxy list
+    checkList = request.json.get('check')
+    proxyList = request.json.get('proxy')
 
+    logging.debug('API create task -> check = %s | proxy = %s' % (checkList, proxyList))
     tasks = []
     for proxy in proxyList:
-        tasks.append({**proxy, 'check': checkList})
-    checkId = Manager.addUnion(tasks)
-    logging.debug('api return check id %s' % checkId)
-
+        tasks.append({
+            **proxy,
+            'check': checkList  # load check items
+        })
+    checkId = Manager.addUnion(tasks)  # add into manager -> get id
+    logging.debug('API return task id -> %s' % checkId)
     return jsonResponse({
         'success': True,
         'id': checkId,
@@ -89,13 +75,10 @@ def createTask() -> Response:
 @webApi.route('/task/<taskId>', methods = ['GET'])
 def getTaskInfo(taskId: str) -> Response:
     if not tokenCheck():  # token check
-        return tokenError()
-    logging.critical('API get task %s info' % taskId)
+        return genError('Invalid token')
+    logging.debug('API get task -> %s' % taskId)
     if not Manager.isUnion(taskId):
-        return jsonResponse({
-            'success': False,
-            'message': 'task id not found',
-        })
+        return genError('Task id not found')
     return jsonResponse({
         'success': True,
         **Manager.getUnion(taskId)
@@ -104,7 +87,7 @@ def getTaskInfo(taskId: str) -> Response:
 
 @webApi.route('/version', methods = ['GET'])
 def getVersion() -> Response:
-    logging.debug('get version -> %s' + Version)
+    logging.debug('API get version -> %s' + Version)
     return jsonResponse({
         'success': True,
         'version': Version,
@@ -113,7 +96,7 @@ def getVersion() -> Response:
 
 def startServer(apiToken: str = '', apiPort: int = 7839) -> None:
     global token
-    token = apiToken
+    token = apiToken  # api token (default empty)
     logging.warning('API server at http://:%i/' % apiPort)
     logging.warning('API ' + ('without token' if apiToken == '' else 'token -> %s' % apiToken))
     pywsgi.WSGIServer(('0.0.0.0', apiPort), webApi).serve_forever()  # powered by gevent

Basis/Compile.py

@@ -7,5 +7,5 @@ from Basis.Logger import logging
 
 def startCompile(dirRange: str = '/') -> None:
     for optimize in [-1, 1, 2]:
-        logging.warning('Python optimize compile -> %s (level = %i)' % (dirRange, optimize))
         compileall.compile_dir(dirRange, quiet = 1, maxlevels = 256, optimize = optimize)
+        logging.warning('Python optimize compile -> %s (level = %i)' % (dirRange, optimize))

Basis/DnsProxy.py

@@ -0,0 +1,37 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+
+import time
+import subprocess
+from Basis.Logger import logging
+
+
+def run(command: list) -> subprocess.Popen:
+    logging.debug('Start dnsproxy -> %s' % command)
+    return subprocess.Popen(command, stdout = subprocess.PIPE, stderr = subprocess.STDOUT)
+
+
+def daemon(process: subprocess.Popen, command: list, gap: int = 2) -> None:  # daemon dnsproxy process
+    while True:  # start daemon
+        time.sleep(gap)  # check time gap
+        if process.poll() is not None:  # unexpected exit
+            logging.warning('dnsproxy unexpected exit')
+            logging.debug('output of dnsproxy\n%s' % process.stdout.read().decode('utf-8'))
+            process = run(command)
+
+
+def start(servers: list or None, port: int = 53, cache: int = 4194304) -> None:  # default cache size -> 4MiB
+    if servers is None or len(servers) == 0:  # use origin dns server
+        logging.info('Skip dnsproxy process')
+        return
+    with open('/etc/resolv.conf', 'w') as dnsConfig:
+        dnsConfig.write('nameserver 127.0.0.1\n')  # system dns settings
+    dnsCommand = [
+        'dnsproxy', '--all-servers',
+        '--port', str(port),
+        '--cache', '--cache-size', str(cache)
+    ]
+    for server in servers:
+        dnsCommand += ['--upstream', server]  # upstream dns servers
+    logging.warning('Start dnsproxy at port %i -> %s' % (port, servers))
+    daemon(run(dnsCommand), dnsCommand)

Basis/Logger.py

@@ -6,8 +6,8 @@ import logging
 from colorlog import ColoredFormatter
 
 logFile = 'runtime.log'
-# logLevel = logging.DEBUG
-logLevel = logging.WARNING
+logLevel = logging.DEBUG
+# logLevel = logging.WARNING
 dateFormat = '%Y-%m-%d %H:%M:%S'
 logFormat = '[%(asctime)s] [%(levelname)s] %(message)s (%(module)s.%(funcName)s)'
 logging.basicConfig(

Basis/Manager.py

@@ -27,10 +27,10 @@ class Task(object):
     def addUnion(self, union: list) -> str:  # add union to manager (include multi tasks)
         tasks = {}  # temporary task storage
         taskIds = []  # task id list for manage union
-        unionId = genFlag(length = 16)  # generate union id (16 bytes)
+        unionId = genFlag(length = 12)  # generate union id (12 bytes)
         logging.debug('Manager start to load union [%s]' % unionId)
         for task in union:
-            taskId = genFlag(length = 24)  # generate task id (24 bytes)
+            taskId = genFlag(length = 16)  # generate task id (16 bytes)
             taskIds.append(taskId)
             tasks[taskId] = {
                 'status': self.__TASK_LOADED,  # task status -> loaded

Basis/Process.py

@@ -194,10 +194,10 @@ class Process(object):
             try:
                 with open(self.__logfile, 'r', encoding = 'utf-8') as fileObject:  # read sub process output
                     self.output = fileObject.read()
-                logging.debug('[%s] Process output -> length = %s' % (self.id, len(self.output)))
+                logging.debug('[%s] Process output capture -> length = %s' % (self.id, len(self.output)))
                 self.__deleteFile(self.__logfile)
             except:
-                logging.error('[%s] Failed to read capture log file -> %s' % (self.id, self.__logfile))
+                logging.error('[%s] Failed to read capture file -> %s' % (self.id, self.__logfile))
         if self.file is not None:  # with config file
             for file in self.file:
                 self.__deleteFile(file['path'])

Builder/__init__.py

@@ -47,7 +47,7 @@ class Builder(object):
     output = None
 
     def __loadClient(self):
-        logging.info('[%s] Builder load %s proxy client at %s -> %s' % (
+        logging.info('[%s] Builder load %s client at %s -> %s' % (
             self.id, self.proxyType,
             'socks5://%s:%i' % (hostFormat(self.socksAddr, v6Bracket = True), self.socksPort), self.proxyInfo
         ))

main.py

@@ -1,64 +1,39 @@
 #!/usr/bin/env python3
 # -*- coding: utf-8 -*-
 
-import os
 import time
 import _thread
-import subprocess
+from Basis import DnsProxy
+from Basis.Check import Check
 from Basis.Logger import logging
+from Basis.Manager import Manager
 from Basis.Api import startServer
 from Basis.Constant import Version
 from Basis.Compile import startCompile
 
-dnsServers = None
-# dnsServers = ['223.5.5.5', '119.28.28.28']
-
-
-def startDnsproxy(command: list) -> subprocess.Popen:
-    logging.debug('start dnsproxy -> %s' % command)
-    return subprocess.Popen(command, stdout = subprocess.PIPE, stderr = subprocess.STDOUT)
+# dnsServers = None
+dnsServers = ['223.5.5.5', '119.28.28.28']
 
 
-def daemonDnsproxy(servers: list or None, port: int = 53, cache: int = 4194304) -> None:  # default cache size -> 4MiB
-    if servers is None or len(servers) == 0:
-        logging.info('skip dnsproxy process')
-        return
-    logging.info('start dnsproxy at port %i -> %s' % (port, servers))
-    os.system('echo "nameserver 127.0.0.1" > /etc/resolv.conf')  # system dns settings
-    dnsCommand = [
-        'dnsproxy', '--all-servers',
-        '--port', str(port),
-        '--cache', '--cache-size', str(cache)
-    ]
-    for server in servers:
-        dnsCommand += ['--upstream', server]
-    dnsproxy = startDnsproxy(dnsCommand)
-    while True:
-        time.sleep(2)  # daemon time gap
-        if dnsproxy.poll() is not None:  # unexpected exit
-            logging.warning('dnsproxy unexpected exit')
-            logging.debug('output of dnsproxy\n%s' % dnsproxy.stdout.read().decode('utf-8'))
-            dnsproxy = startDnsproxy(dnsCommand)
+def runCheck() -> None:  # try to pop a task and check
+    try:
+        taskId, taskInfo = Manager.popTask()
+        logging.warning('[%s] Load new task' % taskId)
+    except:
+        return  # no more task
+    checkResult = Check(taskId, taskInfo)
+    logging.warning('[%s] Task finish' % taskId)
+    Manager.finishTask(taskId, checkResult)
 
 
-from Basis.Check import Check
-from Basis.Manager import Manager
-
 def loopCheck() -> None:
-    while True:
-        time.sleep(2)  # TODO: thread pool working
-        try:
-            taskId, taskInfo = Manager.popTask()
-        except:
-            continue
-        logging.info('new task %s -> %s' % (taskId, taskInfo))
-        ret = Check(taskId, taskInfo)
-        logging.info('check result -> %s' % ret)
-        Manager.finishTask(taskId, ret)
+    while True:  # TODO: thread pool working
+        runCheck()
+        time.sleep(2)
 
 
 logging.warning('ProxyC starts running (%s)' % Version)
 _thread.start_new_thread(startCompile, ('/usr', ))  # python compile (generate .pyc file)
-_thread.start_new_thread(daemonDnsproxy, (dnsServers, 53))  # set system's dns server
+_thread.start_new_thread(DnsProxy.start, (dnsServers, 53))  # start dns server
 _thread.start_new_thread(loopCheck, ())  # start loop check
 startServer(apiToken = '')  # start api server