import subprocess
import time
import pycurl
from io import BytesIO 
import sys

def fetch_proxies():
    command = 'echo "show stat" | sudo socat stdio /var/run/haproxy.sock 2>/dev/null | awk -F, \'$1=="socks5" && !($2~/^(FRONTEND|BACKEND)$/) {print $2,$74}\''
    process = subprocess.Popen(command, stdout=subprocess.PIPE, shell=True)
    output, error = process.communicate()

    result_dict = {}  # Initialize default empty dictionary

    if error is not None:
        return result_dict  # Return the empty dictionary in case of error

    lines = output.decode('utf-8').split('\n')

    for line in lines:
        if line != "":
            key_value = line.split(' ')
            result_dict[key_value[0]] = key_value[1]

    return result_dict


def test_proxy(proxy, url):
    global last_calc_time, download_start_time

    buffer = BytesIO()

    c = pycurl.Curl()
    c.setopt(pycurl.URL, url)
    c.setopt(pycurl.WRITEDATA, buffer)
    c.setopt(pycurl.PROXY, proxy)
    c.setopt(pycurl.PROXYTYPE, pycurl.PROXYTYPE_SOCKS5_HOSTNAME)
    c.setopt(pycurl.NOPROGRESS, False)
    c.setopt(pycurl.XFERINFOFUNCTION, progress)
    c.setopt(pycurl.TIMEOUT, 5)

    download_start_time = time.time()
    last_calc_time = download_start_time

    try:
        c.perform()
    except pycurl.error as e:
        #print(f'An error occurred: {e}')
        # 输出异常的错误代码和描述
        #print(f'Error code: {e.args[0]}')
        #print(f'Error message: {e.args[1]}')
        if e.args[0]==42:
            #print("Script interrupted by user. Exiting During perform pycurl.")
            sys.exit(0)  # 立即退出程序
        else:
            pass

    average_speed = c.getinfo(pycurl.SPEED_DOWNLOAD) / 1024
    return average_speed

def progress(download_t, download_d, upload_t, upload_d):
    global last_calc_time, download_start_time
    current_time = time.time()

    if current_time - last_calc_time >= 2:
        elapsed_time = current_time - download_start_time
        current_speed = download_d / elapsed_time / 1024
        print(f"Current Speed: {current_speed:.2f} kB/s")
        last_calc_time = current_time

proxy_data = fetch_proxies()
url = "http://ipv4.download.thinkbroadband.com/1GB.zip"

for key, value in proxy_data.items():
    print(f"Proxy: {key}")
    try:
        average_speed = test_proxy(value, url)
        print(f"Average Speed: {average_speed:.2f} kB/s")
    except KeyboardInterrupt:
        print("Script interrupted by user. Exiting.")
        break  # 退出for循环