import subprocess
import time
import pycurl
import signal
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 = {}

    if error is not None:
        return result_dict

    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, kb_interrupt

    kb_interrupt = False

    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

    signal.pthread_sigmask(signal.SIG_BLOCK, [signal.SIGINT])

    try:
        c.perform()
    except pycurl.error as e:
        pass

    signal.pthread_sigmask(signal.SIG_UNBLOCK, [signal.SIGINT])

    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, kb_interrupt
    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

    try:
        signal.pthread_sigmask(signal.SIG_UNBLOCK, [signal.SIGINT])
    except KeyboardInterrupt:
        kb_interrupt = True
        print('interrupted!')
        return 1
    finally:
        signal.pthread_sigmask(signal.SIG_BLOCK, [signal.SIGINT])

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

kb_interrupt = False

for key, value in proxy_data.items():
    print(f"Proxy: {key}")

    if kb_interrupt:
        print("Script terminated by user. Exiting.")
        sys.exit(0)

    try:
        average_speed = test_proxy(value, url)
        print(f"Average Speed: {average_speed:.2f} kB/s")
    except KeyboardInterrupt:
        print("Keyboard interrupt encountered during proxy testing. Exiting.")
        break