refresh_ss_config.py

#!/usr/bin/env python3

import os, sys, time, datetime
import argparse

_parser = argparse.ArgumentParser()
_parser.add_argument('-s', '--servers', nargs='+', help='List of servers to test in the format hostname:port.', required=True)
_parser.add_argument('-a', '--servers-alt', nargs='+', help='Specify alternative address for health check corresponding to each item in servers list. Use same format hostname:port.')
_parser.add_argument('-u', '--url', help='The URL to connect for testing. If this parameter is set, the hostname:port in servers list will be used as socks5 proxies. If not set, a socket connection will be attemped directly to hostname:port.')
_parser.add_argument('-c', '--config-name', help='The name of the shadowsocks-libev server configuration.', required=True)
_parser.add_argument('-n', '--count', help='The number of test to perform.', type=int, default=20)
_parser.add_argument('-v', '--verbose', help='Show additional information.', action='store_true')
args = _parser.parse_args()

import socket, subprocess, requests
import multiprocessing as mp

def checkConnection(target):
    host, port = target[0].split(':')
    port = int(port)
    host_check, port_check = host, port

    if target[1]:
        host_check, port_check = target[1].split(':')
        port_check = int(port_check)

    print(f'Checking host {host_check} on port {port_check}...')
    success = 0
    totalTime = 0.0
    for i in range(args.count):
        resp = False
        startTime = datetime.datetime.now()
        if args.url:
            try:
                resp = requests.get(
                    args.url,
                    headers={ 'Connection': 'close' },
                    proxies={ 'http': f'socks5://{host_check}:{port_check}', 'https': f'socks5://{host_check}:{port_check}' }
                ).status_code == 200
            except: pass
        else:
            try:
                sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
                sock.settimeout(1)
                resp = sock.connect_ex((host_check, port_check)) == 0
            except: pass
            finally: sock.close()

        if (resp): success += 1
        
        totalTime += (datetime.datetime.now() - startTime).total_seconds()
        time.sleep(1)
    
    print(f'Total time of {host}:{port} is {totalTime:.2f}s with success rate {success / args.count:.0%}.')
    return ((host, port), success / args.count, totalTime)

if __name__ == '__main__':
    pool = mp.Pool()
    results = pool.map(checkConnection, [(x, args.servers_alt[i] if args.servers_alt else None) for i,x in enumerate(args.servers)])
    results.sort(key=lambda r: (-r[1], r[2]))
    
    if args.verbose: print(results)
    
    server_name, server_port = results[0][0][0], results[0][0][1]
    print(f'Host {server_name}:{server_port} will be used. Updating config...')
    # update config
    subprocess.call(['uci', 'set', f'shadowsocks-libev.{args.config_name}.server={server_name}'])
    subprocess.call(['uci', 'set', f'shadowsocks-libev.{args.config_name}.server_port={server_port}'])
    subprocess.call(['uci', 'commit', 'shadowsocks-libev'])
    subprocess.call(['/etc/init.d/shadowsocks-libev', 'reload'])
    subprocess.call(['logger', '-p', 'info', f'Config {args.config_name} updated with {server_name}:{server_port}.'])