ProxmoxDash/traffic_city.go

package main

import (
	"compress/gzip"
	"crypto/sha1"
	"encoding/hex"
	"fmt"
	"io"
	"net"
	"net/http"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"time"

	"github.com/oschwald/maxminddb-golang"
)

const cityDownloadURL = "https://download.db-ip.com/free/dbip-city-lite-2026-10.mmdb.gz"
const cityDatabaseSHA1 = "df17bd24390108ed1d24dcb7718dc60f6b9abd4f"

var trafficCity struct {
	sync.Mutex
	once    sync.Once
	db      *maxminddb.Reader
	running bool
	message string
}

func trafficCityPath() string {
	db := os.Getenv("DASHBOARD_DB_PATH")
	if db == "" {
		db = defaultDatabasePath
	}
	return filepath.Join(filepath.Dir(db), "dbip-city.mmdb")
}
func trafficCityReady() bool {
	trafficCity.once.Do(func() {
		trafficCity.Lock()
		defer trafficCity.Unlock()
		trafficCity.db, _ = maxminddb.Open(trafficCityPath())
		if trafficCity.db != nil && !strings.Contains(strings.ToLower(trafficCity.db.Metadata.DatabaseType), "city") {
			trafficCity.db.Close()
			trafficCity.db = nil
		}
	})
	trafficCity.Lock()
	defer trafficCity.Unlock()
	return trafficCity.db != nil
}
func lookupTrafficCity(text string) (string, string) {
	ip := net.ParseIP(text)
	if ip == nil || ip.IsPrivate() || ip.IsLoopback() || ip.IsLinkLocalUnicast() || !ip.IsGlobalUnicast() || !trafficCityReady() {
		return "", ""
	}
	trafficCity.Lock()
	defer trafficCity.Unlock()
	var r struct {
		City struct {
			Names map[string]string `maxminddb:"names"`
		} `maxminddb:"city"`
		Subdivisions []struct {
			Names map[string]string `maxminddb:"names"`
		} `maxminddb:"subdivisions"`
	}
	if trafficCity.db.Lookup(ip, &r) != nil {
		return "", ""
	}
	region := ""
	if len(r.Subdivisions) > 0 {
		region = cityName(r.Subdivisions[0].Names)
	}
	return cityName(r.City.Names), region
}
func trafficCityRoutes(mux *http.ServeMux) {
	mux.HandleFunc("/api/settings/traffic/city-database", func(w http.ResponseWriter, r *http.Request) {
		if r.Method != "GET" && r.Method != "POST" {
			w.WriteHeader(405)
			return
		}
		trafficCityReady()
		trafficCity.Lock()
		if r.Method == "POST" && !trafficCity.running && trafficCity.db == nil {
			trafficCity.running = true
			trafficCity.message = "Загрузка базы городов (около 58 МБ)…"
			go func() {
				err := installTrafficCity()
				trafficCity.Lock()
				defer trafficCity.Unlock()
				trafficCity.running = false
				if err != nil {
					trafficCity.message = "Не удалось установить базу: " + err.Error()
				} else {
					trafficCity.message = "База городов DB-IP City Lite 2026-10 установлена"
				}
			}()
		}
		status := map[string]any{"ready": trafficCity.db != nil, "running": trafficCity.running, "message": trafficCity.message, "date": "2026-10"}
		trafficCity.Unlock()
		writeJSON(w, status, nil)
	})
}
func cityMirrorURL() string {
	base := strings.TrimRight(os.Getenv("GITEA_URL"), "/")
	if base == "" {
		base = "https://git.myown.center"
	}
	return base + "/maxim/ProxmoxDash/releases/download/v0.25.1/dbip-city-lite-2026-10.mmdb.gz"
}
func installTrafficCity() error {
	if err := installTrafficCityFrom(cityMirrorURL(), cityDatabaseSHA1); err == nil {
		return nil
	}
	return installTrafficCityFrom(cityDownloadURL, cityDatabaseSHA1)
}
func installTrafficCityFrom(source, expectedHash string) error {
	transport := http.DefaultTransport.(*http.Transport).Clone()
	transport.ResponseHeaderTimeout = 20 * time.Second
	transport.DialContext = (&net.Dialer{Timeout: 10 * time.Second, KeepAlive: 30 * time.Second}).DialContext
	defer transport.CloseIdleConnections()
	client := &http.Client{Timeout: 15 * time.Minute, Transport: transport}
	response, err := client.Get(source)
	if err != nil {
		return err
	}
	defer response.Body.Close()
	if response.StatusCode != 200 {
		return fmt.Errorf("HTTP %d", response.StatusCode)
	}
	reader, err := gzip.NewReader(io.LimitReader(response.Body, 70<<20))
	if err != nil {
		return err
	}
	defer reader.Close()
	path := trafficCityPath()
	if err = os.MkdirAll(filepath.Dir(path), 0750); err != nil {
		return err
	}
	file, err := os.CreateTemp(filepath.Dir(path), ".city-*.mmdb")
	if err != nil {
		return err
	}
	defer os.Remove(file.Name())
	defer file.Close()
	hash := sha1.New()
	n, err := io.Copy(io.MultiWriter(file, hash), io.LimitReader(reader, (160<<20)+1))
	if err != nil {
		return err
	}
	if n > 160<<20 {
		return fmt.Errorf("база превышает лимит размера")
	}
	if hex.EncodeToString(hash.Sum(nil)) != expectedHash {
		return fmt.Errorf("контрольная сумма базы не совпала")
	}
	if err = file.Sync(); err != nil {
		return err
	}
	if err = file.Close(); err != nil {
		return err
	}
	db, err := maxminddb.Open(file.Name())
	if err != nil {
		return err
	}
	if !strings.Contains(strings.ToLower(db.Metadata.DatabaseType), "city") {
		db.Close()
		return fmt.Errorf("неверный тип GeoIP-базы")
	}
	if err = os.Rename(file.Name(), path); err != nil {
		db.Close()
		return err
	}
	trafficCity.Lock()
	defer trafficCity.Unlock()
	if trafficCity.db != nil {
		trafficCity.db.Close()
	}
	trafficCity.db = db
	return nil
}