ProxmoxDash/traffic_city.go

package main

import (
	"compress/gzip"
	"crypto/sha1"
	"encoding/hex"
	"fmt"
	"io"
	"math"
	"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
	cache   map[string]TrafficCityLocation
}

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
}

type TrafficCityLocation struct {
	City      string  `json:"city"`
	Region    string  `json:"region,omitempty"`
	Latitude  float64 `json:"latitude"`
	Longitude float64 `json:"longitude"`
	Located   bool    `json:"located"`
}

func lookupTrafficCity(text string) (string, string) {
	g := lookupTrafficCityLocation(text)
	return g.City, g.Region
}
func lookupTrafficCityLocation(text string) TrafficCityLocation {
	ip := net.ParseIP(text)
	if ip == nil || ip.IsPrivate() || ip.IsLoopback() || ip.IsLinkLocalUnicast() || !ip.IsGlobalUnicast() || !trafficCityReady() {
		return TrafficCityLocation{}
	}
	trafficCity.Lock()
	defer trafficCity.Unlock()
	if g, ok := trafficCity.cache[text]; ok {
		return g
	}
	var r struct {
		Location struct {
			Latitude  *float64 `maxminddb:"latitude"`
			Longitude *float64 `maxminddb:"longitude"`
		} `maxminddb:"location"`
		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 TrafficCityLocation{}
	}
	region := ""
	if len(r.Subdivisions) > 0 {
		region = cityName(r.Subdivisions[0].Names)
	}
	g := TrafficCityLocation{City: cityName(r.City.Names), Region: region}
	if g.City != "" && r.Location.Latitude != nil && r.Location.Longitude != nil {
		g.Latitude, g.Longitude = *r.Location.Latitude, *r.Location.Longitude
		g.Located = !math.IsNaN(g.Latitude) && !math.IsNaN(g.Longitude) && math.Abs(g.Latitude) <= 90 && math.Abs(g.Longitude) <= 180
	}
	if trafficCity.cache == nil {
		trafficCity.cache = map[string]TrafficCityLocation{}
	}
	if len(trafficCity.cache) < 32768 {
		trafficCity.cache[text] = g
	}
	return g
}
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
	trafficCity.cache = nil
	return nil
}