From b543d0c1f7025ff9de9884f069e7d41cfe03c300 Mon Sep 17 00:00:00 2001
From: Guanran Wang <guanran928@outlook.com>
Date: Sun, 20 Oct 2024 16:56:23 +0800
Subject: [PATCH] feat: add asn and city info

---
 main.go | 77 +++++++++++++++++++++++++++++++++++++++++----------------
 1 file changed, 56 insertions(+), 21 deletions(-)

diff --git a/main.go b/main.go
index ba2a97e..efcbc3a 100644
--- a/main.go
+++ b/main.go
@@ -14,30 +14,41 @@ import (
 )
 
 var (
-	db      *geoip2.Reader
+	ASNDB   *geoip2.Reader
+	CityDB  *geoip2.Reader
 	tmpl    *template.Template
 	config  ServerConfig
 	ipCache sync.Map
 )
 
 type ServerConfig struct {
-	listen    string
-	countrydb string
+	Listen string
+	CityDB string
+	ASNDB  string
 }
 
 func main() {
 	config = ServerConfig{
-		listen:    getEnvOr("IP_CHECKER_LISTEN", ":8080"),
-		countrydb: os.Getenv("IP_CHECKER_COUNTRY_DB"),
+		Listen: getEnvOr("IP_CHECKER_LISTEN", ":8080"),
+		CityDB: os.Getenv("IP_CHECKER_CITY_DB"),
+		ASNDB:  os.Getenv("IP_CHECKER_ASN_DB"),
 	}
 
-	if config.countrydb != "" {
+	if config.CityDB != "" {
 		var err error
-		db, err = geoip2.Open(config.countrydb)
+		CityDB, err = geoip2.Open(config.CityDB)
 		if err != nil {
-			log.Fatalf("Error opening GeoIP database: %v", err)
+			log.Fatalf("Error opening GeoIP City database: %v", err)
 		}
-		defer db.Close()
+		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"))
@@ -54,12 +65,14 @@ Disallow: /harm/to/self
 `)
 	})
 
-	log.Printf("Starting server on %s", config.listen)
-	log.Fatal(http.ListenAndServe(config.listen, nil))
+	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"
 	}
@@ -73,29 +86,51 @@ func getIPCountry(ip string) string {
 		return "invalid IP"
 	}
 
-	if config.countrydb == "" {
-		log.Printf("Country database not set. Returning 'unknown'")
+	// check if DB exists
+	// FIXME:
+	if config.CityDB == "" && config.ASNDB == "" {
+		log.Printf("GeoIP database not set. Returning 'unknown'")
 		return "unknown"
 	}
 
-	if country, found := ipCache.Load(ip); found {
-		return country.(string)
+	// cache
+	if cached, found := ipCache.Load(ip); found {
+		return cached.(string)
 	}
 
-	record, err := db.Country(parsedIP)
+	// parse DB
+	cityRecord, err := CityDB.City(parsedIP)
 	if err != nil {
-		log.Printf("GeoIP lookup failed for IP: %s, error: %v", ip, err)
-		return "GeoIP lookup failed"
+		log.Printf("GeoIP City lookup failed for IP: %s, error: %v", ip, err)
+		return "City record lookup failed"
 	}
 
-	country, ok := record.Country.Names["en"]
+	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"
+	}
 
-	ipCache.Store(ip, country)
-	return country
+	result := region + ", " + country + ", " + ASN
+	ipCache.Store(ip, result)
+	return result
 }
 
 func handleRequest(w http.ResponseWriter, r *http.Request) {