From 7f357604ebb6764fbcf3469180096183a19b0050 Mon Sep 17 00:00:00 2001 From: dnomd343 Date: Sun, 7 Aug 2022 20:34:13 +0800 Subject: [PATCH] feat: argparse for input args --- Basis/Api.py | 29 +++++++++++------------- Basis/Constant.py | 12 +++++++--- main.py | 57 ++++++++++++++++++++++++++++++++++++++++++----- 3 files changed, 74 insertions(+), 24 deletions(-) diff --git a/Basis/Api.py b/Basis/Api.py index 506c5c3..8cb5b36 100644 --- a/Basis/Api.py +++ b/Basis/Api.py @@ -1,16 +1,15 @@ #!/usr/bin/env python3 # -*- coding: utf-8 -*- +import os import json from gevent import pywsgi from Checker import formatCheck from Basis.Logger import logging from Basis.Manager import Manager -from Basis.Constant import Version from flask import Flask, Response, request +from Basis.Constant import ApiPort, ApiPath, ApiToken, Version -token = '' -webPath = '/' # root of api server webApi = Flask(__name__) # init flask server @@ -43,16 +42,16 @@ def genError(message: str) -> Response: def tokenCheck() -> bool: - if token == '': return True # without token check + if ApiToken == '': return True # without token check if request.method == 'GET': - return request.args.get('token') == token + return request.args.get('token') == ApiToken elif request.method == 'POST': - return request.json.get('token') == token + return request.json.get('token') == ApiToken else: return False # polyfill -@webApi.route('/task', methods = ['GET']) +@webApi.route(os.path.join(ApiPath, 'task'), methods = ['GET']) def getTaskList() -> Response: if not tokenCheck(): # token check return genError('Invalid token') @@ -64,7 +63,7 @@ def getTaskList() -> Response: }) -@webApi.route('/task', methods = ['POST']) +@webApi.route(os.path.join(ApiPath, 'task'), methods = ['POST']) def createTask() -> Response: if not tokenCheck(): # token check return genError('Invalid token') @@ -93,7 +92,7 @@ def createTask() -> Response: }) -@webApi.route('/task/', methods = ['GET']) +@webApi.route(os.path.join(ApiPath, 'task/'), methods = ['GET']) def getTaskInfo(taskId: str) -> Response: if not tokenCheck(): # token check return genError('Invalid token') @@ -106,7 +105,7 @@ def getTaskInfo(taskId: str) -> Response: }) -@webApi.route('/version', methods = ['GET']) +@webApi.route(os.path.join(ApiPath, 'version'), methods = ['GET']) def getVersion() -> Response: logging.debug('API get version -> %s' + Version) return jsonResponse({ @@ -115,9 +114,7 @@ def getVersion() -> Response: }) -def startServer(apiToken: str = '', apiPort: int = 7839) -> None: - global token - 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 +def startServer() -> None: + 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/Constant.py b/Basis/Constant.py index b81e32e..971509a 100644 --- a/Basis/Constant.py +++ b/Basis/Constant.py @@ -1,15 +1,21 @@ #!/usr/bin/env python3 # -*- coding: utf-8 -*- +# Global Options Version = 'dev' + +ApiPath = '/' +ApiPort = 7839 ApiToken = '' -# ApiToken = 'dnomd343' + # LogLevel = 'DEBUG' -LogLevel = 'WARNING' +# LogLevel = 'WARNING' +LogLevel = 'INFO' LogFile = 'runtime.log' + +DnsServer = [] WorkDir = '/tmp/ProxyC' ObfsSite = 'www.bing.com' -DnsServer = ['223.5.5.5', '119.29.29.29'] PathEnv = '/bin:/sbin:/usr/bin:/usr/sbin:/usr/local/bin' # Shadowsocks Info diff --git a/main.py b/main.py index 36b3622..b652388 100755 --- a/main.py +++ b/main.py @@ -1,11 +1,58 @@ #!/usr/bin/env python3 # -*- coding: utf-8 -*- - +import sys +import copy import time import _thread +import argparse import compileall from Basis import Constant + + +def mainArgParse(rawArgs: list) -> argparse.Namespace: + mainParser = argparse.ArgumentParser(description = 'Start running API server') + mainParser.add_argument('--log', type = str, default = 'debug', help = 'output log level') + mainParser.add_argument('--port', type = int, default = 7839, help = 'port for running') + mainParser.add_argument('--token', type = str, default = '', help = 'token for api server') + mainParser.add_argument('-v', '--version', help = 'show version', action = 'store_true') + # TODO: api path / dns server + return mainParser.parse_args(rawArgs) + + +def testArgParse(rawArgs: list) -> argparse.Namespace: + testParser = argparse.ArgumentParser(description = 'Test that each function is working properly') + testParser.add_argument('PROTOCOL', type = str, help = 'test protocol name') + testParser.add_argument('-a', '--all', help = 'test extra shadowsocks items', action = 'store_true') + testParser.add_argument('-6', '--ipv6', help = 'test on ipv6 network', action = 'store_true') + testParser.add_argument('--url', type = str, default = 'http://baidu.com', help = 'http request url') + testParser.add_argument('--cert', type = str, help = 'specify the certificate id') + testParser.add_argument('--thread', type = int, default = 16, help = 'thread number in check process') + testParser.add_argument('--select', type = str, nargs = '+', help = 'select id list for test') + return testParser.parse_args(rawArgs) + + +inputArgs = copy.copy(sys.argv) +if len(inputArgs) >= 0: # remove first arg (normally file name) + inputArgs.pop(0) +if len(inputArgs) != 0 and inputArgs[0].lower() == 'test': # test mode + inputArgs.pop(0) # remove `test` + if len(inputArgs) == 0 or inputArgs[0].startswith('-'): # no protocol is specified + inputArgs = ['all'] + inputArgs + testArgs = testArgParse(inputArgs) + print(testArgs) + sys.exit(1) + # TODO: start test process +else: + mainArgs = mainArgParse(inputArgs) + if mainArgs.version: # output version and exit + print('ProxyC version -> %s' % Constant.Version) + sys.exit(0) + Constant.LogLevel = mainArgs.log # overwrite global options + Constant.ApiPort = mainArgs.port + Constant.ApiToken = mainArgs.token + + from Basis.Check import Check from Basis import Api, DnsProxy from Basis.Logger import logging @@ -25,7 +72,7 @@ def runCheck(taskId: str, taskInfo: dict) -> None: Manager.finishTask(taskId, checkResult) # commit check result -def loopCheck(threadNum: int = 16) -> None: +def loop(threadNum: int = 16) -> None: threadPool = ThreadPoolExecutor(max_workers = threadNum) # init thread pool while True: try: @@ -38,7 +85,7 @@ def loopCheck(threadNum: int = 16) -> None: logging.warning('ProxyC starts running (%s)' % Constant.Version) -_thread.start_new_thread(pythonCompile, ('/usr', )) # python compile (generate .pyc file) +_thread.start_new_thread(pythonCompile, ('/usr',)) # python compile (generate .pyc file) _thread.start_new_thread(DnsProxy.start, (Constant.DnsServer, 53)) # start dns server -_thread.start_new_thread(loopCheck, ()) # start loop check -Api.startServer(apiToken = Constant.ApiToken) # start api server +_thread.start_new_thread(loop, ()) # start check loop +Api.startServer() # start api server