#!/usr/bin/env python3 # -*- coding: utf-8 -*- import os import time import requests from threading import Thread from Utils.Logger import logger from Utils.Constant import WorkDir, TestHost, TestSite from Utils.Common import md5Sum, genFlag, hostFormat, isVacantPort Settings = { 'workDir': WorkDir, 'site': TestSite, 'serverBind': '', 'clientBind': '', 'host': '', 'cert': '', 'key': '', } def loadBind(serverV6: bool = False, clientV6: bool = False) -> None: Settings['serverBind'] = '::1' if serverV6 else '127.0.0.1' Settings['clientBind'] = '::1' if clientV6 else '127.0.0.1' def waitPort(port: int, times: int = 50, delay: int = 100) -> bool: # wait until port occupied for i in range(times): if not isVacantPort(port): # port occupied return True time.sleep(delay / 1000) # default wait 100ms * 50 => 5s return False # timeout def genCert(host: str, certInfo: dict, remark: str = 'ProxyC') -> None: # generate self-signed certificate certOrgInfo = ['--organization', remark, '--organizationUnit', remark] # organization info logger.critical('Load self-signed certificate') os.system('mkdir -p %s' % Settings['workDir']) # make sure that work directory exist # create CA data at first (by mad) logger.critical('Creating CA certificate and key...') os.system(' '.join( # generate CA certificate and privkey ['mad', 'ca', '--ca', certInfo['caCert'], '--key', certInfo['caKey'], '--commonName', remark] + certOrgInfo )) # generate private key and sign certificate logger.critical('Signing certificate...') os.system(' '.join(['mad', 'cert', '--domain', host] + [ # generate certificate and privkey, then signed by CA '--ca', certInfo['caCert'], '--ca_key', certInfo['caKey'], '--cert', certInfo['cert'], '--key', certInfo['key'], ] + certOrgInfo)) # install CA certificate and record self-signed cert info logger.critical('Installing CA certificate...') os.system('cat %s >> /etc/ssl/certs/ca-certificates.crt' % certInfo['caCert']) # add into system's trust list def loadCert(host: str = TestHost, certId: str = '') -> None: # load certificate newCert = (certId == '') certId = genFlag(length = 8) if certId == '' else certId certInfo = { 'caCert': os.path.join(Settings['workDir'], 'proxyc_%s_ca.pem' % certId), 'caKey': os.path.join(Settings['workDir'], 'proxyc_%s_ca_key.pem' % certId), 'cert': os.path.join(Settings['workDir'], 'proxyc_%s_cert.pem' % certId), 'key': os.path.join(Settings['workDir'], 'proxyc_%s_cert_key.pem' % certId), } if newCert: genCert(host, certInfo) # generate new certificate Settings['host'] = host Settings['cert'] = certInfo['cert'] Settings['key'] = certInfo['key'] logger.warning('Certificate load complete -> ID = %s' % certId) def httpCheck(socksInfo: dict, url: str, testId: str, timeout: int = 10) -> None: socksProxy = 'socks5://%s:%i' % (hostFormat(socksInfo['addr'], v6Bracket = True), socksInfo['port']) try: logger.debug('[%s] Http request via %s' % (testId, socksProxy)) request = requests.get(url, timeout = timeout, proxies = { # http request via socks5 'http': socksProxy, 'https': socksProxy, }) request.raise_for_status() # throw error when server return 4xx or 5xx (don't actually need) logger.info('[%s] %s -> ok' % (testId, socksProxy)) except Exception as exp: logger.error('[%s] %s -> error\n%s' % (testId, socksProxy, exp)) # show detail of error reason raise RuntimeError('Http request via socks5 failed') def runTest(testInfo: dict, testUrl: str, testSelect: set or None, delay: int = 1) -> None: testId = md5Sum(testInfo['caption'])[:12] # generate test ID if testSelect is not None: # testSelect is None -> run all test if testId not in testSelect: # skip unselected task return logger.warning('[%s] %s' % (testId, testInfo['caption'])) # show caption logger.debug('[%s] Server ID -> %s | Client ID -> %s' % ( testId, testInfo['server'].id, testInfo['client'].id )) testInfo['server'].id = testId + '-server' testInfo['client'].id = testId + '-client' # build server and client and wait them start testInfo['server'].start() # start test server if waitPort(testInfo['interface']['port']): # wait for server logger.debug('[%s] Test server start complete' % testId) testInfo['client'].start() # start test client if waitPort(testInfo['socks']['port']): # wait for client logger.debug('[%s] Test client start complete' % testId) # start test process try: logger.debug('[%s] Test process start' % testId) time.sleep(delay) # delay a short time before check httpCheck(testInfo['socks'], testUrl, testId) # run http request test testInfo['client'].quit() # clean up client testInfo['server'].quit() # clean up server except: testInfo['client'].quit() testInfo['server'].quit() logger.warning('[%s] Client info' % testId) logger.error('[%(id)s-server]\n⬤ CMD => %(cmd)s\n⬤ ENV => %(env)s\n⬤ FILE => %(file)s\n%(output)s' % { 'id': testId, 'cmd': testInfo['client'].cmd, 'env': testInfo['client'].env, 'file': testInfo['client'].file, 'output': '-' * 96 + '\n' + testInfo['client'].output + '-' * 96, }) logger.warning('[%s] Server info' % testId) logger.error('[%(id)s-client]\n⬤ CMD => %(cmd)s\n⬤ ENV => %(env)s\n⬤ FILE => %(file)s\n%(output)s' % { 'id': testId, 'cmd': testInfo['server'].cmd, 'env': testInfo['server'].env, 'file': testInfo['server'].file, 'output': '-' * 96 + '\n' + testInfo['server'].output + '-' * 96, }) def Test(testIter: iter, threadNum: int, testUrl: str, testFilter: set or None = None) -> None: threads = [] while True: # infinite loop try: for thread in threads: if thread.is_alive(): # skip running thread continue threads.remove(thread) # remove dead thread if len(threads) < threadNum: for i in range(threadNum - len(threads)): # start threads within limit thread = Thread( # create new thread target = runTest, args = (next(testIter), testUrl, testFilter) ) thread.start() threads.append(thread) # record thread info time.sleep(0.1) except StopIteration: # traverse completed break for thread in threads: # wait until all threads exit thread.join()