Check all mirror of KaOS

kmirror.go

package main

import (
	"bufio"
	"crypto/md5"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"os"
	"strings"
	"sync"
)

const (
	main_mirror = "http://kaosx.tk/repo/$repo"
	mirrorlist  = "/etc/pacman.d/mirrorlist"
	pacconf     = "/etc/pacman.conf"
)

type RepoState struct {
	Name   string
	Md5    string
	Synced bool
}

type MirrorState struct {
	Name   string
	Online bool
	Repos  []*RepoState
}

type CountryState struct {
	Country string
	Mirrors []*MirrorState
}

func getMirrors() (mirrors map[string][]string, err error) {
	f, err := os.Open(mirrorlist)
	if err != nil {
		return
	}
	defer f.Close()
	buf := bufio.NewReader(f)
	var country string
	mirrors = make(map[string][]string)
	for {
		s, e := buf.ReadString('\n')
		if i := strings.Index(s, "Server ="); i >= 0 {
			m := strings.TrimSpace(s[i+8:])
			mirrors[country] = append(mirrors[country], m)
		} else if strings.HasPrefix(s, "#") {
			country = strings.TrimSpace(s[1:])
		}
		if e != nil {
			break
		}
	}
	return
}

func getRepos() (repos []string, err error) {
	f, err := os.Open(pacconf)
	if err != nil {
		return
	}
	defer f.Close()
	buf := bufio.NewReader(f)
	for {
		s, e := buf.ReadString('\n')
		s = strings.TrimSpace(s)
		l := len(s)
		if l > 2 && s[0] == '[' && s[l-1] == ']' && s[1:l-1] != "options" {
			repos = append(repos, s[1:l-1])
		}
		if e != nil {
			break
		}
	}
	return
}

func isMirrorOnline(url string) bool {
	resp, err := http.Head(url)
	return err == nil && resp.StatusCode == 200
}

func getMd5(mirror, repo string) (md5sum string, err error) {
	url := mirror + fmt.Sprintf("%s/%s.db.tar.gz", repo, repo)
	resp, err := http.Get(url)
	defer resp.Body.Close()
	if err == nil {
		if resp.StatusCode != 200 {
			err = fmt.Errorf("[%i] %s", resp.StatusCode, resp.Status)
		} else {
			var b []byte
			if b, err = ioutil.ReadAll(resp.Body); err == nil {
				md5sum = fmt.Sprintf("%x", md5.Sum(b))
			}
		}
	}
	return
}

func updateMirrorState(mirror *MirrorState) {
	mirror.Name = strings.Replace(mirror.Name, "$repo", "", 1)
	if mirror.Online = isMirrorOnline(mirror.Name); !mirror.Online {
		return
	}
	var wg sync.WaitGroup
	for _, r := range mirror.Repos {
		wg.Add(1)
		go func(repo *RepoState) {
			defer wg.Done()
			if md5, err := getMd5(mirror.Name, repo.Name); err == nil {
				repo.Md5 = md5
			}
		}(r)
	}
	wg.Wait()
}

func checkMd5(mirror, mainmirror *MirrorState) {
	if !mirror.Online {
		return
	}
	for i, r := range mainmirror.Repos {
		mirror.Repos[i].Synced = mirror.Repos[i].Md5 != "" && mirror.Repos[i].Md5 == r.Md5
	}
}

func initState(mirrors map[string][]string, repos []string) (countries []*CountryState, mm *MirrorState) {
	for c, cm := range mirrors {
		cs := &CountryState{
			Country: c,
			Mirrors: make([]*MirrorState, len(cm)),
		}
		for i, m := range cm {
			cs.Mirrors[i] = &MirrorState{
				Name:  m,
				Repos: make([]*RepoState, len(repos)),
			}
			for j, r := range repos {
				cs.Mirrors[i].Repos[j] = &RepoState{
					Name: r,
				}
			}
			if m == main_mirror {
				mm = cs.Mirrors[i]
			}
		}
		countries = append(countries, cs)
	}
	return
}

func main() {
	mirrors, err := getMirrors()
	if err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(1)
	}
	repos, err := getRepos()
	if err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(1)
	}
	has_build := false
	for _, r := range repos {
		if has_build = r == "build"; has_build {
			break
		}
	}
	if !has_build {
		repos = append([]string{"build"}, repos...)
	}
	mstate, mm := initState(mirrors, repos)
	var wg sync.WaitGroup
	for _, c := range mstate {
		for _, m := range c.Mirrors {
			wg.Add(1)
			go func(mirror *MirrorState) {
				defer wg.Done()
				updateMirrorState(mirror)
			}(m)
		}
	}
	wg.Wait()
	for _, c := range mstate {
		for _, m := range c.Mirrors {
			checkMd5(m, mm)
		}
	}
	b, _ := json.MarshalIndent(mstate, "", "  ")
	fmt.Println(string(b))
}