feat: add asn and city info

This commit is contained in:
Guanran Wang 2024-10-20 16:56:23 +08:00
parent 51621426b8
commit b543d0c1f7
Signed by: nyancat
GPG key ID: 91F97D9ED12639CF

77
main.go
View file

@ -14,30 +14,41 @@ import (
) )
var ( var (
db *geoip2.Reader ASNDB *geoip2.Reader
CityDB *geoip2.Reader
tmpl *template.Template tmpl *template.Template
config ServerConfig config ServerConfig
ipCache sync.Map ipCache sync.Map
) )
type ServerConfig struct { type ServerConfig struct {
listen string Listen string
countrydb string CityDB string
ASNDB string
} }
func main() { func main() {
config = ServerConfig{ config = ServerConfig{
listen: getEnvOr("IP_CHECKER_LISTEN", ":8080"), Listen: getEnvOr("IP_CHECKER_LISTEN", ":8080"),
countrydb: os.Getenv("IP_CHECKER_COUNTRY_DB"), CityDB: os.Getenv("IP_CHECKER_CITY_DB"),
ASNDB: os.Getenv("IP_CHECKER_ASN_DB"),
} }
if config.countrydb != "" { if config.CityDB != "" {
var err error var err error
db, err = geoip2.Open(config.countrydb) CityDB, err = geoip2.Open(config.CityDB)
if err != nil { 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")) 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.Printf("Starting server on %s", config.Listen)
log.Fatal(http.ListenAndServe(config.listen, nil)) log.Fatal(http.ListenAndServe(config.Listen, nil))
} }
func getIPCountry(ip string) string { func getIPCountry(ip string) string {
parsedIP := net.ParseIP(ip) parsedIP := net.ParseIP(ip)
// check for reserved ip
if parsedIP.IsPrivate() { if parsedIP.IsPrivate() {
return "private IP" return "private IP"
} }
@ -73,29 +86,51 @@ func getIPCountry(ip string) string {
return "invalid IP" return "invalid IP"
} }
if config.countrydb == "" { // check if DB exists
log.Printf("Country database not set. Returning 'unknown'") // FIXME:
if config.CityDB == "" && config.ASNDB == "" {
log.Printf("GeoIP database not set. Returning 'unknown'")
return "unknown" return "unknown"
} }
if country, found := ipCache.Load(ip); found { // cache
return country.(string) 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 { if err != nil {
log.Printf("GeoIP lookup failed for IP: %s, error: %v", ip, err) log.Printf("GeoIP City lookup failed for IP: %s, error: %v", ip, err)
return "GeoIP lookup failed" 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 { if !ok {
log.Printf("Country name not found in GeoIP data for IP: %s", ip) log.Printf("Country name not found in GeoIP data for IP: %s", ip)
return "unknown" 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) result := region + ", " + country + ", " + ASN
return country ipCache.Store(ip, result)
return result
} }
func handleRequest(w http.ResponseWriter, r *http.Request) { func handleRequest(w http.ResponseWriter, r *http.Request) {