Prism/internal/geodata/geodata.go

package geodata

import (
	"fmt"
	"io"
	"net/http"
	"os"
	"path/filepath"
	"strings"
	"time"

	"github.com/metacubex/mihomo/component/geodata"
	C "github.com/metacubex/mihomo/constant"
)

const DefaultBase = "https://ghfast.top/https://github.com/MetaCubeX/meta-rules-dat/releases/download/latest"

type Config struct {
	CountryMmdbURL string
	GeoIPURL       string
	GeoSiteURL     string
}

type FileSpec struct {
	Name string
	URL  func(Config) string
}

var fileSpecs = []FileSpec{
	{Name: "Country.mmdb", URL: func(c Config) string { return c.CountryMmdbURL }},
	{Name: "geoip.dat", URL: func(c Config) string { return c.GeoIPURL }},
	{Name: "geosite.dat", URL: func(c Config) string { return c.GeoSiteURL }},
}

func DefaultConfig() Config {
	return Config{
		CountryMmdbURL: DefaultBase + "/Country.mmdb",
		GeoIPURL:       DefaultBase + "/geoip.dat",
		GeoSiteURL:     DefaultBase + "/geosite.dat",
	}
}

func ConfigFromSettings(settings map[string]string) Config {
	def := DefaultConfig()
	return Config{
		CountryMmdbURL: pickURL(settings["geo_country_mmdb_url"], def.CountryMmdbURL),
		GeoIPURL:       pickURL(settings["geo_geoip_url"], def.GeoIPURL),
		GeoSiteURL:     pickURL(settings["geo_geosite_url"], def.GeoSiteURL),
	}
}

func pickURL(value, fallback string) string {
	value = strings.TrimSpace(value)
	if value == "" {
		return fallback
	}
	return value
}

func InitHomeDir(dataDir string) string {
	geoDir := filepath.Join(dataDir, "geo")
	_ = os.MkdirAll(geoDir, 0o755)
	C.SetHomeDir(geoDir)
	return geoDir
}

func ApplyMihomoURLs(cfg Config) {
	geodata.SetMmdbUrl(cfg.CountryMmdbURL)
	geodata.SetGeoIpUrl(cfg.GeoIPURL)
	geodata.SetGeoSiteUrl(cfg.GeoSiteURL)
	geodata.SetASNUrl(DefaultBase + "/GeoLite2-ASN.mmdb")
}

type FileStatus struct {
	Name   string `json:"name"`
	Exists bool   `json:"exists"`
	Size   int64  `json:"size"`
	URL    string `json:"url"`
}

func Status(dataDir string, cfg Config) []FileStatus {
	geoDir := filepath.Join(dataDir, "geo")
	out := make([]FileStatus, 0, len(fileSpecs))
	for _, spec := range fileSpecs {
		dest := filepath.Join(geoDir, spec.Name)
		item := FileStatus{Name: spec.Name, URL: spec.URL(cfg)}
		if st, err := os.Stat(dest); err == nil && st.Size() > 0 {
			item.Exists = true
			item.Size = st.Size()
		}
		out = append(out, item)
	}
	return out
}

func Ready(dataDir string) bool {
	geoDir := filepath.Join(dataDir, "geo")
	for _, spec := range fileSpecs {
		st, err := os.Stat(filepath.Join(geoDir, spec.Name))
		if err != nil || st.Size() == 0 {
			return false
		}
	}
	return true
}

// Ensure downloads missing geo files. When force is true, all files are re-downloaded.
func Ensure(dataDir string, cfg Config, force bool) ([]string, error) {
	geoDir := InitHomeDir(dataDir)
	client := &http.Client{Timeout: 120 * time.Second}
	var updated []string
	for _, spec := range fileSpecs {
		dest := filepath.Join(geoDir, spec.Name)
		if !force {
			if st, err := os.Stat(dest); err == nil && st.Size() > 0 {
				continue
			}
		}
		url := spec.URL(cfg)
		if err := downloadFile(client, url, dest); err != nil {
			return updated, fmt.Errorf("download %s: %w (place files manually in %s)", spec.Name, err, geoDir)
		}
		updated = append(updated, spec.Name)
	}
	return updated, nil
}

func downloadFile(client *http.Client, url, dest string) error {
	resp, err := client.Get(url)
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	if resp.StatusCode >= 400 {
		return fmt.Errorf("HTTP %d", resp.StatusCode)
	}
	tmp := dest + ".tmp"
	f, err := os.Create(tmp)
	if err != nil {
		return err
	}
	if _, err := io.Copy(f, resp.Body); err != nil {
		f.Close()
		os.Remove(tmp)
		return err
	}
	f.Close()
	return os.Rename(tmp, dest)
}