zurg/pkg/realdebrid/network.go

package realdebrid

import (
	"bufio"
	"fmt"
	"net/http"
	"os/exec"
	"sort"
	"strconv"
	"strings"
	"sync"
	"time"
)

type IPInfo struct {
	Address string
	Hops    int
	Latency time.Duration
}

func traceroute(ip string) (int, time.Duration, error) {
	// Try executing traceroute
	cmd := exec.Command("traceroute", "-n", "-q", "1", "-w", "1", ip)
	out, err := cmd.CombinedOutput()

	if err == nil {
		// Traceroute executed successfully
		output := string(out)
		lines := strings.Split(output, "\n")
		hopCount := len(lines) - 1
		var latency time.Duration

		if hopCount > 0 {
			lastLine := lines[hopCount-1]
			fmt.Print(".")
			parts := strings.Fields(lastLine)
			if len(parts) >= 3 {
				latencyValue, parseErr := strconv.ParseFloat(parts[2], 64)
				if parseErr == nil {
					latency = time.Duration(latencyValue * float64(time.Millisecond))
				}
			}
		}
		return hopCount, latency, nil
	}

	// Traceroute not successful, measure latency using ping
	pingCmd := exec.Command("ping", "-c", "1", ip)
	pingOut, pingErr := pingCmd.CombinedOutput()

	if pingErr != nil {
		return 0, 0, pingErr
	}

	fmt.Print(".")
	pingOutput := string(pingOut)
	pingLines := strings.Split(pingOutput, "\n")
	for _, line := range pingLines {
		if strings.Contains(line, "time=") {
			parts := strings.Split(line, " ")
			for _, part := range parts {
				if strings.Contains(part, "time=") {
					timeStr := strings.Split(part, "=")[1]
					timeStr = strings.TrimSuffix(timeStr, " ms")
					latencyValue, parseErr := strconv.ParseFloat(timeStr, 64)
					if parseErr == nil {
						latency := time.Duration(latencyValue * float64(time.Millisecond))
						return -1, latency, nil
					}
				}
			}
		}
	}

	return 0, 0, fmt.Errorf("failed to measure latency")
}

func fetchIPs(url string) ([]string, error) {
	resp, err := http.Get(url)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()

	var ips []string
	scanner := bufio.NewScanner(resp.Body)
	for scanner.Scan() {
		ips = append(ips, scanner.Text())
	}
	if err := scanner.Err(); err != nil {
		return nil, err
	}
	return ips, nil
}

func RunTest(netTestType string) {
	fmt.Print("Running network test...")

	ipv4URL := "https://gist.githubusercontent.com/yowmamasita/d0c1c7353500d0928cb5242484e8ed06/raw/ipv4.txt"
	ipv6URL := "https://gist.githubusercontent.com/yowmamasita/d0c1c7353500d0928cb5242484e8ed06/raw/ipv6.txt"

	var ips []string

	if netTestType == "ipv4" || netTestType == "both" {
		ipv4IPs, err := fetchIPs(ipv4URL)
		if err != nil {
			fmt.Println("Error fetching IPv4 IPs:", err)
			return
		}
		ips = append(ips, ipv4IPs...)
	}

	if netTestType == "ipv6" || netTestType == "both" {
		ipv6IPs, err := fetchIPs(ipv6URL)
		if err != nil {
			fmt.Println("Error fetching IPv6 IPs:", err)
			return
		}
		ips = append(ips, ipv6IPs...)
	}

	var wg sync.WaitGroup
	infoChan := make(chan IPInfo, len(ips))
	semaphore := make(chan struct{}, 10)

	for _, ip := range ips {
		wg.Add(1)
		semaphore <- struct{}{}
		go func(ip string) {
			defer wg.Done()
			hops, latency, err := traceroute(ip)
			if err != nil {
				fmt.Printf("Error performing traceroute for %s:%s\n", ip, err)
			} else {
				infoChan <- IPInfo{Address: ip, Hops: hops, Latency: latency}
			}
			<-semaphore
		}(ip)
	}

	wg.Wait()
	close(semaphore)
	close(infoChan)
	fmt.Printf("complete!\n\n")

	var ipInfos []IPInfo
	for info := range infoChan {
		ipInfos = append(ipInfos, info)
	}

	sort.Slice(ipInfos, func(i, j int) bool {
		return ipInfos[i].Latency < ipInfos[j].Latency
	})

	const minResults = 10
	const maxResults = 20
	var okIPs []IPInfo

	if len(ipInfos) > 0 {
		// Start by adding the best IPs based on hops and latency up to the minResults
		for i := 0; i < min(len(ipInfos), minResults); i++ {
			okIPs = append(okIPs, ipInfos[i])
		}

		// Find the highest latency in the current okIPs list
		highestLatency := okIPs[len(okIPs)-1].Latency

		// Add any additional IPs that have latency within a reasonable range of the highest latency
		for _, info := range ipInfos[minResults:] {
			if len(okIPs) >= maxResults {
				break // Stop adding IPs if maxResults is reached
			}
			if info.Latency <= highestLatency+(highestLatency/3) {
				okIPs = append(okIPs, info)
			}
		}
	}

	fmt.Printf("Here are the results, you can copy-paste the following to your config.yml:\n\n")
	fmt.Println("preferred_hosts:")
	for _, info := range okIPs {
		fmt.Printf("  - %s # hops: %d latency: %v\n", info.Address, info.Hops, info.Latency)
	}
}