first version for testing

fastmirror.go

@@ -1,14 +1,19 @@
 package fastmirror
 
 import (
+	"bufio"
 	"fmt"
+	"io"
 	"io/ioutil"
 	"log"
+	"net/http"
 	"os"
 	"os/exec"
 	"path/filepath"
 	"runtime"
+	"strconv"
 	"strings"
+	"time"
 )
 
 func CLIEntry() {
@@ -19,54 +24,80 @@ func CLIEntry() {
 		log.Fatal("This program is only for Ubuntu")
 	}
 
-	fp := identifySourcesFile()
+	sourcesFilePath := findSourcesFilePath()
 
-	fmt.Printf("Found sources file: %s\n", fp)
+	fmt.Printf("Found sources file: %s\n", sourcesFilePath)
 
-	architecture := runtime.GOARCH
+	// Create backup of the sources file
-	switch architecture {
+	createBackup(sourcesFilePath)
-	case "amd64", "386":
+
-		fmt.Println("Use archive.ubuntu.com for packages")
+	// Create sources.list.d directory if it doesn't exist
-	default:
+	createSourcesListD()
-		fmt.Println("Use ports.ubuntu.com for packages")
+
-	}
+	// 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.Fatal("Failed to run lsb_release: %v\n", err)
+		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.Fatal("Failed to run lsb_release: %v\n", err)
+		log.Fatalf("Failed to run lsb_release: %v\n", err)
 	}
 
 	distro := strings.TrimSpace(string(output))
-	if distro == "Ubuntu" {
+	return distro == "Ubuntu"
-		return true
-	} else {
-		return false
-	}
 }
 
-func identifySourcesFile() string {
+func findSourcesFilePath() string {
-	const sourcesList = "/etc/apt/sources.list"
+	const sourcesListPath = "/etc/apt/sources.list"
-	const sourcesListD = "/etc/apt/sources.list.d/"
+	const sourcesListDPath = "/etc/apt/sources.list.d/"
-	if checkFileForUbuntu(sourcesList) {
+	if isUbuntuSourcesFile(sourcesListPath) {
-		fmt.Printf("Found Ubuntu sources file: %s\n", sourcesList)
+		fmt.Printf("Found Ubuntu sources file: %s\n", sourcesListPath)
-		return sourcesList
+		return sourcesListPath
 	}
-	files, err := ioutil.ReadDir(sourcesListD)
+	files, err := ioutil.ReadDir(sourcesListDPath)
 	if err != nil {
-		fmt.Printf("Failed to read directory %s: %v\n", sourcesListD, err)
+		fmt.Printf("Failed to read directory %s: %v\n", sourcesListDPath, err)
 		os.Exit(1)
 	}
 
@@ -75,8 +106,8 @@ func identifySourcesFile() string {
 			continue
 		}
 
-		filePath := filepath.Join(sourcesListD, file.Name())
+		filePath := filepath.Join(sourcesListDPath, file.Name())
-		if checkFileForUbuntu(filePath) {
+		if isUbuntuSourcesFile(filePath) {
 			fmt.Printf("Found Ubuntu sources file: %s\n", filePath)
 			return filePath
 		}
@@ -84,16 +115,158 @@ func identifySourcesFile() string {
 	return ""
 }
 
-func checkFileForUbuntu(filePath string) bool {
+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
 	}
 
-	if strings.Contains(strings.ToLower(string(content)), "ubuntu.com") {
+	return strings.Contains(strings.ToLower(string(content)), "ubuntu.com")
-		return true
 }
 
+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)
+}