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

def sigint_handler(signal, frame):
    print(f"KeyboardInterrupt at line {frame.f_lineno}")
    sys.exit(0)

signal.signal(signal.SIGINT, sigint_handler)

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}\''
    try:
        process = subprocess.Popen(command, stdout=subprocess.PIPE, shell=True)
        output, error = process.communicate()
    except Exception as e:
        print("Error occurred while executing shell command: ", e)
        return {}

    result_dict = {}

    if error is not None:
        print("Error occurred while getting proxies status: ", error.decode('utf-8'))

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

    for line in lines:
        if line and len(line.split(' ')) == 2:
            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, total_downloaded_data
    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
    total_downloaded_data = 0
    py_speed = 0

    try:
        c.perform()
        total_downloaded_data = c.getinfo(pycurl.SIZE_DOWNLOAD)
    except pycurl.error as e:
        print("Error occurred while performing curl: ", e)
    finally:
        py_speed = c.getinfo(pycurl.SPEED_DOWNLOAD) / 1024
        c.close()

    elapsed_time = time.time() - download_start_time
    average_speed = total_downloaded_data / elapsed_time / 1024
    print(f"PycURL Speed: {py_speed:.2f} kB/s")

    return average_speed


def progress(download_t, download_d, upload_t, upload_d):
    global last_calc_time, download_start_time, total_downloaded_data
    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

    total_downloaded_data = download_d

    return 0

proxy_data = fetch_proxies()

if not proxy_data:
    print("No proxy data found.")
    sys.exit(1)

url = "http://ipv4.download.thinkbroadband.com/1GB.zip"

for key, value in proxy_data.items():
    print(f"Testing Proxy: {key}")
    average_speed = test_proxy(value, url)
    print(f"Average Speed: {average_speed:.2f} kB/s")
