From d1f76b91c66b09cecfa29ce178b178a6d2729c6a Mon Sep 17 00:00:00 2001 From: dnomd343 Date: Tue, 2 Aug 2022 16:33:21 +0800 Subject: [PATCH] update: optimize code --- Basis/Api.py | 65 +++++++++++++++++---------------------------- Basis/Compile.py | 2 +- Basis/DnsProxy.py | 37 ++++++++++++++++++++++++++ Basis/Logger.py | 4 +-- Basis/Manager.py | 4 +-- Basis/Process.py | 4 +-- Builder/__init__.py | 2 +- main.py | 61 +++++++++++++----------------------------- 8 files changed, 87 insertions(+), 92 deletions(-) create mode 100644 Basis/DnsProxy.py diff --git a/Basis/Api.py b/Basis/Api.py index c577417..1222425 100644 --- a/Basis/Api.py +++ b/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/', 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 diff --git a/Basis/Compile.py b/Basis/Compile.py index 22aea3d..ad90c81 100644 --- a/Basis/Compile.py +++ b/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)) diff --git a/Basis/DnsProxy.py b/Basis/DnsProxy.py new file mode 100644 index 0000000..3b738f9 --- /dev/null +++ b/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) diff --git a/Basis/Logger.py b/Basis/Logger.py index 764bf9e..383cee2 100644 --- a/Basis/Logger.py +++ b/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( diff --git a/Basis/Manager.py b/Basis/Manager.py index 6360ce2..e9154d3 100644 --- a/Basis/Manager.py +++ b/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 diff --git a/Basis/Process.py b/Basis/Process.py index c2f134a..3e3b2a0 100644 --- a/Basis/Process.py +++ b/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']) diff --git a/Builder/__init__.py b/Builder/__init__.py index a0de0e4..dd1a60e 100644 --- a/Builder/__init__.py +++ b/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 )) diff --git a/main.py b/main.py index d9bcbe0..d6990db 100755 --- a/main.py +++ b/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