fastmirror/sources.go

package fastmirror

import (
	"bufio"
	_ "embed"
	"fmt"
	"html/template"
	"log"
	"os"
	"path/filepath"
	"strings"
)

//go:embed sources_list_legacy.tmpl
var legacyTemplate string

//go:embed sources_list_modern.tmpl
var modernTemplate string

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

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

		filePath := filepath.Join(sourcesListDPath, file.Name())
		if isUbuntuSourcesFile(filePath) {
			log.Printf("Found Ubuntu sources file: %s", filePath)
			return filePath, nil
		}
	}
	return "", fmt.Errorf("no Ubuntu sources file found")
}

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

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

func createBackup(filePath string) error {
	backupPath := filePath + ".bak"
	if _, err := os.Stat(backupPath); err == nil {
		return fmt.Errorf("backup file %s already exists", backupPath)
	}
	if err := os.Rename(filePath, backupPath); err != nil {
		return fmt.Errorf("failed to create backup of %s: %v", filePath, err)
	}
	log.Printf("Created backup: %s", backupPath)
	return nil
}

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

func extractSuites(filePath string) ([]string, error) {
	content, err := os.ReadFile(filePath)
	if err != nil {
		return nil, fmt.Errorf("failed to read file %s: %v", filePath, err)
	}

	var suites []string
	scanner := bufio.NewScanner(strings.NewReader(string(content)))

	isModernFormat := false
	for scanner.Scan() {
		line := scanner.Text()
		if strings.HasPrefix(line, "Types:") {
			isModernFormat = true
			break
		}
	}

	scanner = bufio.NewScanner(strings.NewReader(string(content)))
	if isModernFormat {
		for scanner.Scan() {
			line := scanner.Text()
			if strings.HasPrefix(line, "Suites:") {
				parts := strings.Fields(line)
				suites = append(suites, parts[1:]...)
			}
		}
	} else {
		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 {
		return nil, fmt.Errorf("failed to scan file %s: %v", filePath, err)
	}
	return suites, nil
}

func updateSourcesFile(filePath, mirrorURL string, suites []string) error {
	codename, err := getUbuntuCodename()
	if err != nil {
		return err
	}
	majorVersion, minorVersion, err := getUbuntuVersion()
	if err != nil {
		return err
	}

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

	var content string
	data := struct {
		MirrorURL string
		Codename  string
		Suites    string
	}{
		MirrorURL: mirrorURL,
		Codename:  codename,
		Suites:    suitesString,
	}

	if majorVersion > 24 || (majorVersion == 24 && minorVersion >= 4) {
		tmpl, err := template.New("modern").Parse(modernTemplate)
		if err != nil {
			return fmt.Errorf("failed to parse modern template: %v", err)
		}
		var buf strings.Builder
		err = tmpl.Execute(&buf, data)
		if err != nil {
			return fmt.Errorf("failed to execute modern template: %v", err)
		}
		content = buf.String()
	} else {
		tmpl, err := template.New("legacy").Parse(legacyTemplate)
		if err != nil {
			return fmt.Errorf("failed to parse legacy template: %v", err)
		}
		var buf strings.Builder
		err = tmpl.Execute(&buf, data)
		if err != nil {
			return fmt.Errorf("failed to execute legacy template: %v", err)
		}
		content = buf.String()
	}

	err = os.WriteFile(filePath, []byte(content), 0644)
	if err != nil {
		return fmt.Errorf("failed to update sources file %s: %v", filePath, err)
	}
	log.Printf("Updated sources file: %s", filePath)
	return nil
}