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 (
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) {