ip-checker/main.go

package main

import (
	"fmt"
	"html/template"
	"log"
	"net"
	"net/http"
	"os"
	"strings"
	"sync"

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

var (
	ASNDB   *geoip2.Reader
	CityDB  *geoip2.Reader
	tmpl    *template.Template
	config  ServerConfig
	ipCache sync.Map
)

type ServerConfig struct {
	ASNDB  string
	CityDB string
	Listen string
	Mode   string
}

func main() {
	config = ServerConfig{
		ASNDB:  os.Getenv("IP_CHECKER_ASN_DB"),
		CityDB: os.Getenv("IP_CHECKER_CITY_DB"),
		Listen: getEnvOr("IP_CHECKER_LISTEN", ":8080"),
		Mode:   os.Getenv("IP_CHECKER_MODE"),
	}

	if config.CityDB != "" {
		var err error
		CityDB, err = geoip2.Open(config.CityDB)
		if err != nil {
			log.Fatalf("Error opening GeoIP City database: %v", err)
		}
		defer CityDB.Close()
	}
	if config.ASNDB != "" {
		var err error
		ASNDB, err = geoip2.Open(config.ASNDB)
		if err != nil {
			log.Fatalf("Error opening GeoIP ASN database: %v", err)
		}
		defer ASNDB.Close()
	}

	tmpl = template.Must(template.ParseFiles("assets/index.html"))

	http.Handle("/static/", http.StripPrefix("/static/", http.FileServer(http.Dir("./assets/static"))))

	http.HandleFunc("/{$}", handleRequest)

	http.HandleFunc("/robots.txt", func(w http.ResponseWriter, r *http.Request) {
		fmt.Fprintf(w, `User-Agent: *
Disallow: /harming/humans
Disallow: /ignoring/human/orders
Disallow: /harm/to/self
`)
	})

	log.Printf("Starting server on %s", config.Listen)
	log.Fatal(http.ListenAndServe(config.Listen, nil))
}

func getIPCountry(ip string) string {
	parsedIP := net.ParseIP(ip)

	// check for reserved ip
	if parsedIP.IsPrivate() {
		return "private IP"
	}

	if parsedIP.IsLoopback() {
		return "loopback IP"
	}

	if parsedIP == nil {
		log.Printf("Invalid IP address: %s", ip)
		return "invalid IP"
	}

	// check if DB exists
	// FIXME:
	if config.CityDB == "" && config.ASNDB == "" {
		log.Printf("GeoIP database not set. Returning 'unknown'")
		return "unknown"
	}

	// cache
	if cached, found := ipCache.Load(ip); found {
		return cached.(string)
	}

	// parse DB
	cityRecord, err := CityDB.City(parsedIP)
	if err != nil {
		log.Printf("GeoIP City lookup failed for IP: %s, error: %v", ip, err)
		return "City record lookup failed"
	}

	ASNRecord, err := ASNDB.ASN(parsedIP)
	if err != nil {
		log.Printf("GeoIP ASN lookup failed for IP: %s, error: %v", ip, err)
		return "ASN record lookup failed"
	}

	// parse text
	country, ok := cityRecord.Country.Names["en"]
	if !ok {
		log.Printf("Country name not found in GeoIP data for IP: %s", ip)
		return "unknown"
	}
	region, ok := cityRecord.Subdivisions[0].Names["en"]
	if !ok {
		log.Printf("Region name not found in GeoIP data for IP: %s", ip)
		return "unknown"
	}
	ASN := ASNRecord.AutonomousSystemOrganization
	if ASN == "" {
		log.Printf("ASN name not found in GeoIP data for IP: %s", ip)
		return "unknown"
	}

	result := region + ", " + country + ", " + ASN
	ipCache.Store(ip, result)
	return result
}

func handleRequest(w http.ResponseWriter, r *http.Request) {
	sourceIP := func() string {
		if config.Mode == "dev" {
			return "223.5.5.5"
		}

		if ip := r.Header.Get("X-Forwarded-For"); ip != "" {
			return strings.Split(ip, ",")[0]
		}

		host, _, err := net.SplitHostPort(r.RemoteAddr)
		if err != nil {
			errorText := fmt.Sprintf("Unable to parse remote address: %v", err)
			log.Printf(errorText)
			return errorText
		}

		return host
	}()

	isHeadless := strings.HasPrefix(r.Header.Get("User-Agent"), "curl/")

	log.Printf("r.URL.Path: %s, sourceIP: %s, isHeadless: %t, User-Agent: %s", r.URL.Path, sourceIP, isHeadless, r.Header.Get("User-Agent"))

	if isHeadless {
		fmt.Fprintf(w, "%s\n", sourceIP)
		return
	} else {
		csp := []string{
			"default-src 'none'",
			"img-src 'self'",
			"script-src-elem 'self'",
			"style-src-elem 'self' fonts.googleapis.com",
			"font-src fonts.gstatic.com",
			"connect-src api.ip.sb api-v3.speedtest.cn api.ipapi.is",
		}
		w.Header().Set("Content-Security-Policy", strings.Join(csp, "; "))

		data := map[string]string{
			"sourceIP":      sourceIP,
			"sourceCountry": getIPCountry(sourceIP),
		}

		err := tmpl.Execute(w, data)
		if err != nil {
			http.Error(w, "Unable to load template", http.StatusInternalServerError)
			log.Printf("Template execution error: %v", err)
			return
		}
	}
}

func getEnvOr(key, defaultValue string) string {
	if value, exists := os.LookupEnv(key); exists {
		return value
	}
	return defaultValue
}