From a0980a4c678dc958499e2c6c47bd77f95151056c Mon Sep 17 00:00:00 2001 From: sneak Date: Wed, 22 May 2024 09:14:11 -0700 Subject: [PATCH] first version for testing --- fastmirror.go | 233 +++++++++++++++++++++++++++++++++++++++++++------- 1 file changed, 203 insertions(+), 30 deletions(-) diff --git a/fastmirror.go b/fastmirror.go index ed3058e..0e021ff 100644 --- a/fastmirror.go +++ b/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 - switch architecture { - case "amd64", "386": - fmt.Println("Use archive.ubuntu.com for packages") - default: - fmt.Println("Use ports.ubuntu.com for packages") - } + // 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.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 true - } else { - return false - } + return distro == "Ubuntu" } -func identifySourcesFile() string { - const sourcesList = "/etc/apt/sources.list" - const sourcesListD = "/etc/apt/sources.list.d/" - if checkFileForUbuntu(sourcesList) { - fmt.Printf("Found Ubuntu sources file: %s\n", sourcesList) - return sourcesList +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(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()) - if checkFileForUbuntu(filePath) { + filePath := filepath.Join(sourcesListDPath, file.Name()) + 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 true + 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) } - return false + 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) }