fastmirror/fastmirror.go

package fastmirror

import (
	"bufio"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"os/exec"
	"path/filepath"
	"runtime"
	"strconv"
	"strings"
	"time"
)

func CLIEntry() {
	if runtime.GOOS != "linux" {
		log.Fatal("This program is only for Linux")
	}
	if !isUbuntu() {
		log.Fatal("This program is only for Ubuntu")
	}

	sourcesFilePath := findSourcesFilePath()

	fmt.Printf("Found sources file: %s\n", sourcesFilePath)

	// Create backup of the sources file
	createBackup(sourcesFilePath)

	// Create sources.list.d directory if it doesn't exist
	createSourcesListD()

	// Fetch and find the fastest mirror
	fastestMirrorURL := findFastestMirror()
	fmt.Printf("Fastest mirror: %s\n", fastestMirrorURL)

	// Extract suites from the existing sources file
	suites := extractSuites(sourcesFilePath)

	// Update sources file with the fastest mirror
	updateSourcesFile(sourcesFilePath, fastestMirrorURL, suites)
}

func getUbuntuCodename() string {
	cmd := exec.Command("lsb_release", "-cs")
	output, err := cmd.Output()
	if err != nil {
		log.Fatalf("Failed to run lsb_release: %v\n", err)
	}
	codename := strings.TrimSpace(string(output))
	return codename
}

func getUbuntuVersion() (int, int) {
	cmd := exec.Command("lsb_release", "-rs")
	output, err := cmd.Output()
	if err != nil {
		log.Fatalf("Failed to run lsb_release: %v\n", err)
	}
	version := strings.TrimSpace(string(output))
	parts := strings.Split(version, ".")
	if len(parts) != 2 {
		log.Fatalf("Unexpected version format: %s", version)
	}
	majorVersion, err := strconv.Atoi(parts[0])
	if err != nil {
		log.Fatalf("Failed to parse major version: %v", err)
	}
	minorVersion, err := strconv.Atoi(parts[1])
	if err != nil {
		log.Fatalf("Failed to parse minor version: %v", err)
	}
	return majorVersion, minorVersion
}

func isUbuntu() bool {
	cmd := exec.Command("lsb_release", "-is")
	output, err := cmd.Output()
	if err != nil {
		log.Fatalf("Failed to run lsb_release: %v\n", err)
	}

	distro := strings.TrimSpace(string(output))
	return distro == "Ubuntu"
}

func findSourcesFilePath() string {
	const sourcesListPath = "/etc/apt/sources.list"
	const sourcesListDPath = "/etc/apt/sources.list.d/"
	if isUbuntuSourcesFile(sourcesListPath) {
		fmt.Printf("Found Ubuntu sources file: %s\n", sourcesListPath)
		return sourcesListPath
	}
	files, err := ioutil.ReadDir(sourcesListDPath)
	if err != nil {
		fmt.Printf("Failed to read directory %s: %v\n", sourcesListDPath, err)
		os.Exit(1)
	}

	for _, file := range files {
		if file.IsDir() {
			continue
		}

		filePath := filepath.Join(sourcesListDPath, file.Name())
		if isUbuntuSourcesFile(filePath) {
			fmt.Printf("Found Ubuntu sources file: %s\n", filePath)
			return filePath
		}
	}
	return ""
}

func isUbuntuSourcesFile(filePath string) bool {
	content, err := ioutil.ReadFile(filePath)
	if err != nil {
		fmt.Printf("Failed to read file %s: %v\n", filePath, err)
		return false
	}

	return strings.Contains(strings.ToLower(string(content)), "ubuntu.com")
}

func createBackup(filePath string) {
	backupPath := filePath + ".bak"
	err := os.Rename(filePath, backupPath)
	if err != nil {
		log.Fatalf("Failed to create backup of %s: %v\n", filePath, err)
	}
	fmt.Printf("Created backup: %s\n", backupPath)
}

func createSourcesListD() {
	const sourcesListDPath = "/etc/apt/sources.list.d/"
	if _, err := os.Stat(sourcesListDPath); os.IsNotExist(err) {
		err = os.Mkdir(sourcesListDPath, 0755)
		if err != nil {
			log.Fatalf("Failed to create directory %s: %v\n", sourcesListDPath, err)
		}
		fmt.Printf("Created directory: %s\n", sourcesListDPath)
	}
}

func findFastestMirror() string {
	resp, err := http.Get("http://mirrors.ubuntu.com/mirrors.txt")
	if err != nil {
		log.Fatalf("Failed to fetch mirrors list: %v\n", err)
	}
	defer resp.Body.Close()

	body, err := io.ReadAll(resp.Body)
	if err != nil {
		log.Fatalf("Failed to read mirrors list: %v\n", err)
	}

	mirrors := strings.Split(string(body), "\n")
	var fastestMirrorURL string
	var fastestTime time.Duration
	codename := getUbuntuCodename()
	architecture := runtime.GOARCH

	httpClient := http.Client{
		Timeout: 1 * time.Second,
	}

	for _, mirror := range mirrors {
		if strings.HasPrefix(mirror, "https://") {
			mirror = strings.TrimSuffix(mirror, "/")
			if isValidMirror(httpClient, mirror, codename, architecture) {
				startTime := time.Now()
				resp, err := httpClient.Get(mirror)
				if err == nil {
					elapsedTime := time.Since(startTime)
					resp.Body.Close()
					if fastestMirrorURL == "" || elapsedTime < fastestTime {
						fastestMirrorURL = mirror
						fastestTime = elapsedTime
					}
				}
			}
		}
	}

	if fastestMirrorURL == "" {
		log.Fatal("No suitable HTTPS mirror found")
	}

	return fastestMirrorURL
}

func isValidMirror(httpClient http.Client, mirrorURL, codename, architecture string) bool {
	uri := fmt.Sprintf("%s/dists/%s/Release", mirrorURL, codename)
	resp, err := httpClient.Get(uri)
	if err != nil {
		return false
	}
	defer resp.Body.Close()
	return resp.StatusCode == http.StatusOK
}

func extractSuites(filePath string) []string {
	content, err := ioutil.ReadFile(filePath)
	if err != nil {
		log.Fatalf("Failed to read file %s: %v\n", filePath, err)
	}

	// Extract suites from the sources file
	suites := []string{}
	scanner := bufio.NewScanner(strings.NewReader(string(content)))
	for scanner.Scan() {
		line := scanner.Text()
		if strings.HasPrefix(line, "deb ") || strings.HasPrefix(line, "deb-src ") {
			parts := strings.Fields(line)
			if len(parts) >= 4 {
				suites = append(suites, parts[3:]...)
			}
		}
	}
	if err := scanner.Err(); err != nil {
		log.Fatalf("Failed to scan file %s: %v\n", filePath, err)
	}
	return suites
}

func updateSourcesFile(filePath, mirrorURL string, suites []string) {
	architecture := runtime.GOARCH
	codename := getUbuntuCodename()
	majorVersion, minorVersion := getUbuntuVersion()

	var uriPath string
	if architecture == "amd64" || architecture == "386" {
		uriPath = "/ubuntu"
	} else {
		uriPath = "/ubuntu-ports"
	}

	suitesString := strings.Join(suites, " ")

	var content string
	if majorVersion > 24 || (majorVersion == 24 && minorVersion >= 4) {
		content = fmt.Sprintf(`Types: deb
URIs: %s%s
Suites: %s
Components: %s
Signed-By: /usr/share/keyrings/ubuntu-archive-keyring.gpg

## Ubuntu security updates. Aside from URIs and Suites,
## this should mirror your choices in the previous section.
Types: deb
URIs: %s%s
Suites: %s-security
Components: %s
Signed-By: /usr/share/keyrings/ubuntu-archive-keyring.gpg
`, mirrorURL, uriPath, codename, suitesString, mirrorURL, uriPath, codename, suitesString)
	} else {
		content = fmt.Sprintf(`deb %s%s %s %s
deb %s%s %s-updates %s
deb %s%s %s-backports %s
deb %s%s %s-security %s
`, mirrorURL, uriPath, codename, suitesString, mirrorURL, uriPath, codename, suitesString, mirrorURL, uriPath, codename, suitesString, mirrorURL, uriPath, codename, suitesString)
	}

	err := ioutil.WriteFile(filePath, []byte(content), 0644)
	if err != nil {
		log.Fatalf("Failed to update sources file %s: %v\n", filePath, err)
	}
	fmt.Printf("Updated sources file: %s\n", filePath)
}