whois-geoip-web/main.go

//-----------------------------
//Distributed under GPL
//Author Sergey Kalinin
//svk@nuk-svk.ru
//------------------------------

package main

import (
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"html/template"
	"io"
	"log"
	"net"
	"net/http"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"time"

	"github.com/go-co-op/gocron"
	"github.com/gorilla/mux"
	"github.com/likexian/whois"
	"github.com/oschwald/geoip2-golang"
	"golang.org/x/text/language"
)

type GeoData struct {
	// IP адрес
	IPAddress string `json:"ip_address"`
	// Местоположение
	Location struct {
		Country    string  `json:"country,omitempty"`
		City       string  `json:"city,omitempty"`
		Continent  string  `json:"continent,omitempty"`
		Region     string  `json:"region,omitempty"`
		RegionISO  string  `json:"region_iso,omitempty"`
		Latitude   float64 `json:"latitude,omitempty"`
		Longitude  float64 `json:"longitude,omitempty"`
		TimeZone   string  `json:"time_zone,omitempty"`
		PostalCode string  `json:"postal_code,omitempty"`
	} `json:"location"`
	// Данные о сети
	Network struct {
		ASN   uint   `json:"asn,omitempty"`
		ASOrg string `json:"autonomous_system_organization,omitempty"`
	} `json:"network"`
	// ISO коды
	Codes struct {
		CountryISO   string `json:"country,omitempty"`
		ContinentISO string `json:"continent,omitempty"`
	} `json:"codes"`
	// Прокси
	Security struct {
		IsAnonymousProxy    bool `json:"is_anonymous_proxy,omitempty"`
		IsSatelliteProvider bool `json:"is_satellite_provider,omitempty"`
	} `json:"security"`
	// Локаль
	Metadata struct {
		LocaleUsed string `json:"locale_used,omitempty"`
		DataSource string `json:"data_source,omitempty"`
	} `json:"metadata"`
}

// TemplateData структура для передачи данных в шаблон
type TemplateData struct {
	*GeoData
	Error     string
	JSONData  string
	WhoIsData string
	BaseURL   string
	DBVersion string
}

type Config struct {
	MMDBURL          string
	MMDBLocalPath    string
	HTMLTemplatePath string
	ListenPort       string
}

var (
	cityDB         *geoip2.Reader
	countryDB      *geoip2.Reader
	asnDB          *geoip2.Reader
	templates      *template.Template
	cfg            *Config
	dbMutex        sync.RWMutex
	dbFilesVersion string
)

func main() {
	var err error
	// Загрузка конфигурации
	cfg = LoadConfig()

	// Запускаем скачивание файлов
	mmdbDownload("GeoLite2-City.mmdb")
	mmdbDownload("GeoLite2-Country.mmdb")
	mmdbDownload("GeoLite2-ASN.mmdb")

	// Запускаем планировщик параллельно с основным процессом
	go sheduler()

	// Открываем БД
	pathGeoLite2City := filepath.Join(cfg.MMDBLocalPath, "GeoLite2-City.mmdb")
	log.Println(pathGeoLite2City)

	cityDB, err = geoip2.Open(pathGeoLite2City)
	if err != nil {
		log.Fatal("Failed to open City database:", err)
	}
	defer cityDB.Close()

	pathGeoLite2Country := filepath.Join(cfg.MMDBLocalPath, "GeoLite2-Country.mmdb")
	countryDB, err = geoip2.Open(pathGeoLite2Country)
	if err != nil {
		log.Fatal("Failed to open Country database:", err)
	}
	defer countryDB.Close()

	pathGeoLite2ASN := filepath.Join(cfg.MMDBLocalPath, "GeoLite2-ASN.mmdb")
	asnDB, err = geoip2.Open(pathGeoLite2ASN)
	if err != nil {
		log.Fatal("Failed to open ASN database:", err)
	}
	defer asnDB.Close()

	dbFilesVersion = getFileModTime(pathGeoLite2City)

	// Загружаем шаблон
	templates = template.Must(template.ParseFiles(filepath.Join(cfg.HTMLTemplatePath, "index.html")))

	// Обработка ссылок
	r := mux.NewRouter()
	r.HandleFunc("/", homeHandler)
	r.HandleFunc("/api/{ip}", apiHandler)
	r.HandleFunc("/api/", apiCurrentHandler)
	r.HandleFunc("/lookup/{ip}", lookupHandler)
	r.HandleFunc("/{ip}", lookupHandler)

	log.Println("Server starting on:", cfg.ListenPort)
	log.Fatal(http.ListenAndServe(":"+cfg.ListenPort, r))
}

func getFileModTime(filePath string) string {
	// Получаем информацию о файле
	fileInfo, err := os.Stat(filePath)
	if err != nil {
		log.Println("Ошибка в getFileModTime:", err)
		return "Error"
	}

	// Получаем время последнего изменения
	modTime := fileInfo.ModTime()
	// fmt.Println("Время изменения файла:", modTime)

	// Форматируем вывод
	// fmt.Printf("Форматированное время: %s\n", modTime.Format("02/01/2006 15:04:05"))
	return modTime.Format("02/01/2006 15:04:05")
}

// reopenDBs закрывает и заново открывает нужную БД
func reopenDBs(fileName string) {
	dbMutex.Lock()
	defer dbMutex.Unlock()

	var err error

	switch fileName {
	case "GeoLite2-City.mmdb":
		if cityDB != nil {
			cityDB.Close()
		}
		// Открываем базы данных заново
		pathGeoLite2City := filepath.Join(cfg.MMDBLocalPath, fileName)
		cityDB, err = geoip2.Open(pathGeoLite2City)
		if err != nil {
			log.Printf("Ошибка открытия базы данных %v: %v", fileName, err)
			return
		}
	case "GeoLite2-Country.mmdb":
		if countryDB != nil {
			countryDB.Close()
		}
		pathGeoLite2Country := filepath.Join(cfg.MMDBLocalPath, fileName)
		countryDB, err = geoip2.Open(pathGeoLite2Country)
		if err != nil {
			log.Printf("Ошибка открытия базы данных %v: %v", fileName, err)
			return
		}
	case "GeoLite2-ASN.mmdb":
		if asnDB != nil {
			asnDB.Close()
		}
		pathGeoLite2ASN := filepath.Join(cfg.MMDBLocalPath, fileName)
		asnDB, err = geoip2.Open(pathGeoLite2ASN)
		if err != nil {
			log.Printf("Ошибка открытия базы данных %v: %v", fileName, err)
			return
		}
	}
	log.Println("Открыта база данных:", fileName)
	dbFilesVersion = getFileModTime(filepath.Join(cfg.MMDBLocalPath, fileName))
}

func homeHandler(w http.ResponseWriter, r *http.Request) {
	baseURL := getBaseURL(r)

	// Обрабатываем параметр ip из GET запроса
	ipStr := r.URL.Query().Get("ip")
	var geoData *GeoData
	var errMsg string
	var whoisData string

	if ipStr != "" {
		ip := net.ParseIP(ipStr)
		if ip == nil {
			errMsg = "Неверный IP адрес: " + ipStr
		} else {
			preferredLocale := getPreferredLocale(r)
			var err error
			geoData, err = getGeoData(ip, preferredLocale)
			if err != nil {
				errMsg = err.Error()
			}
			whoisData = getWhoIsInfo(ipStr)
		}
	}

	data := TemplateData{
		GeoData:   geoData,
		WhoIsData: whoisData,
		Error:     errMsg,
		BaseURL:   baseURL,
		DBVersion: dbFilesVersion,
	}

	err := templates.ExecuteTemplate(w, "index.html", data)
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
	}
}

func apiCurrentHandler(w http.ResponseWriter, r *http.Request) {
	// Определяем IP клиента
	ipStr := getClientIP(r)
	ip := net.ParseIP(ipStr)
	if ip == nil {
		sendJSONError(w, "Unable to determine client IP", http.StatusBadRequest)
		return
	}

	// Определяем локаль
	preferredLocale := getPreferredLocale(r)
	geoData, err := getGeoData(ip, preferredLocale)
	if err != nil {
		sendJSONError(w, err.Error(), http.StatusInternalServerError)
		return
	}

	sendJSONResponse(w, geoData)
}

func apiHandler(w http.ResponseWriter, r *http.Request) {
	vars := mux.Vars(r)
	ipStr := vars["ip"]

	// Проверяем IP
	ip := net.ParseIP(ipStr)
	if ip == nil {
		sendJSONError(w, "Invalid IP address", http.StatusBadRequest)
		return
	}

	// Определяем локаль
	preferredLocale := getPreferredLocale(r)
	geoData, err := getGeoData(ip, preferredLocale)
	if err != nil {
		sendJSONError(w, err.Error(), http.StatusInternalServerError)
		return
	}

	sendJSONResponse(w, geoData)
}

// Обрабатываем запросы вида /lookup/{ip} и возвращает HTML страницу с результатами
func lookupHandler(w http.ResponseWriter, r *http.Request) {
	vars := mux.Vars(r)
	ipStr := vars["ip"]
	baseURL := getBaseURL(r)

	// Проверка IP
	ip := net.ParseIP(ipStr)
	if ip == nil {
		// Если IP невалидный, показываем ошибку на странице
		templates.ExecuteTemplate(w, "index.html", TemplateData{
			Error:   "Invalid IP address: " + ipStr,
			BaseURL: baseURL,
		})
		return
	}

	// Определяем локаль
	preferredLocale := getPreferredLocale(r)
	geoData, err := getGeoData(ip, preferredLocale)
	if err != nil {
		templates.ExecuteTemplate(w, "index.html", TemplateData{
			Error:   err.Error(),
			BaseURL: baseURL,
		})
		return
	}

	jsonData, err := json.MarshalIndent(geoData, "", "  ")
	if err != nil {
		jsonData = []byte("{\"error\": \"Failed to generate JSON\"}")
	}

	whoisData := getWhoIsInfo(ipStr)

	// Передаем данные в шаблон
	templates.ExecuteTemplate(w, "index.html", TemplateData{
		GeoData:   geoData,
		JSONData:  string(jsonData),
		WhoIsData: whoisData,
		BaseURL:   baseURL,
		DBVersion: dbFilesVersion,
	})
}

// Определяем локаль броузера
func getPreferredLocale(r *http.Request) string {
	acceptLang := r.Header.Get("Accept-Language")
	if acceptLang == "" {
		return "en"
	}

	tags, _, err := language.ParseAcceptLanguage(acceptLang)
	if err != nil || len(tags) == 0 {
		return "en"
	}

	preferred := tags[0]
	base, _ := preferred.Base()

	locale := base.String()

	supportedLocales := []string{"en", "de", "es", "fr", "ja", "pt-BR", "ru", "zh-CN"}

	for _, supported := range supportedLocales {
		if locale == supported {
			return locale
		}
	}

	// Возвращем en если ничего другого не найдено
	return "en"
}

func sendJSONResponse(w http.ResponseWriter, data interface{}) {
	w.Header().Set("Content-Type", "application/json; charset=utf-8")
	w.Header().Set("Access-Control-Allow-Origin", "*")

	// Pretty print JSON with indentation
	jsonData, err := json.MarshalIndent(data, "", "  ")
	if err != nil {
		sendJSONError(w, "Failed to encode JSON response", http.StatusInternalServerError)
		return
	}

	w.Write(jsonData)
}

func sendJSONError(w http.ResponseWriter, message string, statusCode int) {
	w.Header().Set("Content-Type", "application/json; charset=utf-8")
	w.Header().Set("Access-Control-Allow-Origin", "*")
	w.WriteHeader(statusCode)

	errorResponse := map[string]interface{}{
		"error": map[string]string{
			"message": message,
			"code":    http.StatusText(statusCode),
		},
		"success": false,
	}

	jsonData, _ := json.MarshalIndent(errorResponse, "", "  ")
	w.Write(jsonData)
}

// Определяем адрес клиента
func getClientIP(r *http.Request) string {
	if forwarded := r.Header.Get("X-Forwarded-For"); forwarded != "" {
		return forwarded
	}
	if realIP := r.Header.Get("X-Real-IP"); realIP != "" {
		return realIP
	}

	host, _, err := net.SplitHostPort(r.RemoteAddr)
	if err != nil {
		return r.RemoteAddr
	}
	return host
}

// Проверяем доступную локаль
func getLocalizedString(names map[string]string, preferredLocale string) string {
	if names == nil {
		return ""
	}

	// Определяем локаль
	if name, exists := names[preferredLocale]; exists && name != "" {
		return name
	}

	// Переключаем на  en
	if name, exists := names["en"]; exists {
		return name
	}

	// Если нет Английского возвращем пустоту
	for _, name := range names {
		if name != "" {
			return name
		}
	}

	return ""
}

func hasRegisteredCountryData(registeredCountry struct {
	Names             map[string]string `maxminddb:"names"`
	IsoCode           string            `maxminddb:"iso_code"`
	GeoNameID         uint              `maxminddb:"geoname_id"`
	IsInEuropeanUnion bool              `maxminddb:"is_in_european_union"`
}) bool {
	return registeredCountry.IsoCode != "" || len(registeredCountry.Names) > 0
}

// Получаем данные из БД MaxMind
func getGeoData(ip net.IP, preferredLocale string) (*GeoData, error) {
	dbMutex.RLock()
	defer dbMutex.RUnlock()

	data := &GeoData{}
	data.IPAddress = ip.String()
	data.Metadata.LocaleUsed = preferredLocale

	// Получаем данные из БД City
	if city, err := cityDB.City(ip); err == nil {
		data.Metadata.DataSource = "city"

		// Сперва определяем страну (main country)
		data.Location.Country = getLocalizedString(city.Country.Names, preferredLocale)
		data.Codes.CountryISO = city.Country.IsoCode

		// Если (main country) пустое то пробуем (registered country)
		if data.Location.Country == "" && hasRegisteredCountryData(city.RegisteredCountry) {
			data.Location.Country = getLocalizedString(city.RegisteredCountry.Names, preferredLocale)
			if data.Codes.CountryISO == "" {
				data.Codes.CountryISO = city.RegisteredCountry.IsoCode
			}
		}

		// Определяем континет (main continent)
		data.Location.Continent = getLocalizedString(city.Continent.Names, preferredLocale)
		data.Codes.ContinentISO = city.Continent.Code

		// Город и данные местоположения
		data.Location.City = getLocalizedString(city.City.Names, preferredLocale)
		data.Location.Latitude = city.Location.Latitude
		data.Location.Longitude = city.Location.Longitude
		data.Location.TimeZone = city.Location.TimeZone
		data.Location.PostalCode = city.Postal.Code
		data.Security.IsAnonymousProxy = city.Traits.IsAnonymousProxy
		data.Security.IsSatelliteProvider = city.Traits.IsSatelliteProvider

		// Регион
		if len(city.Subdivisions) > 0 {
			subdivision := city.Subdivisions[0]
			data.Location.Region = getLocalizedString(subdivision.Names, preferredLocale)
			data.Location.RegionISO = subdivision.IsoCode
		}
	}

	// Определяем страну (fallback for country data)
	if data.Location.Country == "" {
		if country, err := countryDB.Country(ip); err == nil {
			if data.Metadata.DataSource == "" {
				data.Metadata.DataSource = "country"
			}

			// Сперва определяем страну (main country)
			countryName := getLocalizedString(country.Country.Names, preferredLocale)
			countryISO := country.Country.IsoCode

			// Если (main country) пустое то пробуем (registered country)
			if countryName == "" && hasRegisteredCountryData(country.RegisteredCountry) {
				countryName = getLocalizedString(country.RegisteredCountry.Names, preferredLocale)
				if countryISO == "" {
					countryISO = country.RegisteredCountry.IsoCode
				}
			}

			// Обновлем данные если что нашли
			if countryName != "" {
				data.Location.Country = countryName
			}
			if countryISO != "" {
				data.Codes.CountryISO = countryISO
			}

			// Континент
			continentName := getLocalizedString(country.Continent.Names, preferredLocale)
			continentISO := country.Continent.Code

			if continentName != "" {
				data.Location.Continent = continentName
			}
			if continentISO != "" {
				data.Codes.ContinentISO = continentISO
			}

			// Прокси
			data.Security.IsAnonymousProxy = country.Traits.IsAnonymousProxy
			data.Security.IsSatelliteProvider = country.Traits.IsSatelliteProvider
		}
	}

	// Читаем данные из бд ASN
	if asn, err := asnDB.ASN(ip); err == nil {
		data.Network.ASN = asn.AutonomousSystemNumber
		data.Network.ASOrg = asn.AutonomousSystemOrganization
	}

	return data, nil
}

// Определяем базовый URL сайта где запуцщен проект
func getBaseURL(r *http.Request) string {
	scheme := "http"
	if r.TLS != nil || r.Header.Get("X-Forwarded-Proto") == "https" {
		scheme = "https"
	}

	host := r.Host
	// Убираем порт если это стандартный порт
	if strings.HasSuffix(host, ":80") && scheme == "http" {
		host = strings.TrimSuffix(host, ":80")
	} else if strings.HasSuffix(host, ":443") && scheme == "https" {
		host = strings.TrimSuffix(host, ":443")
	}

	return scheme + "://" + host
}

// Запрос информации с серверов  whois
func getWhoIsInfo(address string) string {
	result, err := whois.Whois(address)

	if err != nil {
		return "WhoIs request error"
	}

	return result
}

// Загрузка конфигурации
func LoadConfig() *Config {
	return &Config{
		MMDBURL:          getEnv("MMDB_URL", "localhost"),
		MMDBLocalPath:    getEnv("MMDB_LOCAL_PATH", "/usr/local/share/geoip/db"),
		HTMLTemplatePath: getEnv("HTML_TEMPLATE_PATH", "/usr/local/share/geoip/templates"),
		ListenPort:       getEnv("LISTEN_PORT", "8080"),
	}
}

// Получение знаячений переменных окружения
func getEnv(key, defaultValue string) string {
	if value := os.Getenv(key); value != "" {
		return value
	}
	return defaultValue
}

// Запуск задания скачивания файлов
func sheduler() {
	// инициализируем объект планировщика
	s := gocron.NewScheduler(time.UTC)
	// добавляем одну задачу на каждую минуту
	s.Cron("30 2 * * *").Do(startDownload)
	// s.Cron("*/1 * * * *").Do(startDownload)
	// запускаем планировщик с блокировкой текущего потока
	s.StartBlocking()
}

// Отслеживает прогресс загрузки
type WriteCounter struct {
	Total uint64
}

// Параллельный запуск загрузки БД
func startDownload() {
	go mmdbDownload("GeoLite2-City.mmdb")
	go mmdbDownload("GeoLite2-Country.mmdb")
	go mmdbDownload("GeoLite2-ASN.mmdb")

}

func (wc *WriteCounter) Write(p []byte) (int, error) {
	n := len(p)
	wc.Total += uint64(n)
	wc.PrintProgress()
	return n, nil
}

func (wc WriteCounter) PrintProgress() {
	fmt.Printf("\rЗагружено %d байт...", wc.Total)
}

// Проверка целостности файла.
func isValidMMDB(filePath string) bool {
	db, err := geoip2.Open(filePath)
	if err != nil {
		return false
	}
	db.Close()
	return true
}

// Проверяем контрольную сумму локального файла
func getFileChecksum(filePath string) (string, error) {
	file, err := os.Open(filePath)
	if err != nil {
		return "", err
	}
	defer file.Close()

	hash := sha256.New()
	if _, err := io.Copy(hash, file); err != nil {
		return "", err
	}

	return hex.EncodeToString(hash.Sum(nil)), nil
}

// Проверяем контрольную сумму удаленного файла (скачиваемого)
func getRemoteChecksum(url string) string {
	resp, err := http.Get(url)
	if err != nil {
		return ""
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		return ""
	}

	body, err := io.ReadAll(resp.Body)
	if err != nil {
		return ""
	}

	// Предполагаем, что файл содержит только хеш
	return strings.TrimSpace(string(body))
}

func shouldDownload(fileName string) bool {
	log.Printf("Проверяем контрольную сумму файла %s", fileName)
	fileURL := cfg.MMDBURL + fileName
	filePath := filepath.Join(cfg.MMDBLocalPath, fileName)

	// Если файла нет, нужно скачивать
	if _, err := os.Stat(filePath); os.IsNotExist(err) {
		return true
	}

	// Проверяем возраст файла (скачиваем если старше 7 дней)
	/*
	   if info, err := os.Stat(filePath); err == nil {
	       if time.Since(info.ModTime()) > 7*24*time.Hour {
	           return true
	       }
	   }

	*/
	// Проверка по контрольной сумме
	if checksumURL := fileURL + ".sha256"; checksumURL != "" {
		remoteFileChecksum := getRemoteChecksum(checksumURL)
		// fmt.Println("Remote", remoteFileChecksum)

		if remoteFileChecksum != "" {
			localFileChecksum, error := getFileChecksum(filePath)
			if error == nil {
				if localFileChecksum != "" {
					fmt.Println(filePath, "Remote:", remoteFileChecksum, "Local:", localFileChecksum)
					fmt.Println(remoteFileChecksum)
					fmt.Println(localFileChecksum)
					if localFileChecksum != remoteFileChecksum {
						return false
					}
				}
			} else {
				fmt.Println("Error", filePath, error)
				return true
			}
		}
	}

	return false
}

// Загрузка файлов БД с ВЭБ-сервера
func mmdbDownload(fileName string) {
	fileURL := cfg.MMDBURL + fileName
	filePath := filepath.Join(cfg.MMDBLocalPath, fileName)

	if !shouldDownload(fileName) {
		log.Printf("Файл %s актуален, пропускаем загрузку", fileName)
		return
	}

	// Скачиваем во временный файл
	tempPath := filePath + ".tmp"

	log.Println("Загружается файл:", fileURL)
	// Создаём файл
	file, err := os.Create(tempPath)
	if err != nil {
		log.Printf("Ошибка создания файла %s: %v", tempPath, err)
		return
	}
	defer file.Close()

	// Загружаем данные с отслеживанием прогресса
	counter := &WriteCounter{}
	response, err := http.Get(fileURL)
	if err != nil {
		log.Printf("Ошибка загрузки файла %s: %v", fileURL, err)
		return
	}
	defer response.Body.Close()

	if response.StatusCode != http.StatusOK {
		log.Printf("Ошибка загрузки: %s", response.Status)
		return
	}

	// Копируем через буфер с отслеживанием прогресса
	_, err = io.Copy(file, io.TeeReader(response.Body, counter))
	if err != nil {
		log.Printf("Ошибка копирования данных: %v", err)
		return
	}

	log.Println("Загрузка завершена! Файл сохранен:", tempPath)

	// Переоткрываем базы данных после успешной загрузки
	// reopenDBs(fileName)
	// Проверяем валидность перед заменой
	if isValidMMDB(tempPath) {
		// Заменяем старый файл
		os.Rename(tempPath, filePath)
		// Переоткрываем только эту базу
		reopenDBs(fileName)
	} else {
		log.Printf("Скачанный файл %s поврежден, удаляем", fileName)
		os.Remove(tempPath)
	}
}