Copy

package geodb

import (
	"context"
	"fmt"
	"honeypot/internal/utils"
	"net"
	"net/netip"
	"strings"

	"io"
	"net/http"
	"os"

	"github.com/oschwald/maxminddb-golang/v2"
)

// ASN minimal struct
type asnLookup struct {
	AutonomousSystemNumber       uint   `maxminddb:"autonomous_system_number" json:"autonomous_system_number"`
	AutonomousSystemOrganization string `maxminddb:"autonomous_system_organization" json:"autonomous_system_organization"`
}

// Minimal "names" struct for obtaining English names
type names struct {
	En string `maxminddb:"en" json:"en"`
}

// Only interested in city and its English name, and containing country info
type cityLookup struct {
	City struct {
		Names names `maxminddb:"names"`
	} `maxminddb:"city"`
	Country struct {
		IsoCode string `maxminddb:"iso_code"`
		Names   names  `maxminddb:"names"`
	} `maxminddb:"country"`
	Location struct {
		Latitude  float64 `maxminddb:"latitude"`
		Longitude float64 `maxminddb:"longitude"`
	} `maxminddb:"location"`
}

type IPMetadata struct {
	IP          string
	Country     string
	CountryCode string
	ASN         int
	ASNOrg      string
	City        string
	Latitude    float64
	Longitude   float64
	FQDN        string
	Domain      string
}

type GeoDB struct {
	ASNDB  *maxminddb.Reader
	CityDB *maxminddb.Reader
}

func NewGeoDB(asnDBPath, cityDBPath string) (*GeoDB, error) {
	if asnDBPath == "" || cityDBPath == "" {
		return nil, fmt.Errorf("database paths cannot be empty")
	}

	asnDb, err := maxminddb.Open(asnDBPath)
	if err != nil {
		return nil, fmt.Errorf("error opening ASN database: %w", err)
	}

	cityDb, err := maxminddb.Open(cityDBPath)
	if err != nil {
		asnDb.Close()
		return nil, fmt.Errorf("error opening city database: %w", err)
	}

	return &GeoDB{
		ASNDB:  asnDb,
		CityDB: cityDb,
	}, nil
}

func (g *GeoDB) Close() error {
	if g.ASNDB != nil {
		if err := g.ASNDB.Close(); err != nil {
			return err
		}
	}
	if g.CityDB != nil {
		if err := g.CityDB.Close(); err != nil {
			return err
		}
	}
	return nil
}

func (g *GeoDB) Reload(asnDBPath, cityDBPath string) error {
	newAsnDb, err := maxminddb.Open(asnDBPath)
	if err != nil {
		return fmt.Errorf("error opening ASN database: %w", err)
	}

	newCityDb, err := maxminddb.Open(cityDBPath)
	if err != nil {
		newAsnDb.Close()
		return fmt.Errorf("error opening city database: %w", err)
	}

	oldAsn := g.ASNDB
	oldCity := g.CityDB

	g.ASNDB = newAsnDb
	g.CityDB = newCityDb

	if oldAsn != nil {
		err := oldAsn.Close()
		if err != nil {
			return err
		}
	}
	if oldCity != nil {
		err := oldCity.Close()
		if err != nil {
			return err
		}
	}

	return nil
}

func DownloadFile(url, filepath string) error {
	// #nosec G107
	resp, err := http.Get(url)
	if err != nil {
		return err
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		return fmt.Errorf("bad status: %s", resp.Status)
	}

	out, err := os.Create(filepath)
	if err != nil {
		return err
	}
	defer out.Close()

	_, err = io.Copy(out, resp.Body)
	return err
}

func (g *GeoDB) LookupMetadata(ctx context.Context, ipStr string) (*IPMetadata, error) {
	parsedIp, err := netip.ParseAddr(ipStr)
	if err != nil {
		return nil, err
	}

	meta := &IPMetadata{
		IP: ipStr,
	}

	// ASN
	var asn asnLookup
	if err := g.ASNDB.Lookup(parsedIp).Decode(&asn); err == nil {
		meta.ASN = int(asn.AutonomousSystemNumber)
		meta.ASNOrg = asn.AutonomousSystemOrganization
	}

	// City/Country/Location
	var city cityLookup
	if err := g.CityDB.Lookup(parsedIp).Decode(&city); err == nil {
		meta.City = city.City.Names.En
		meta.Country = city.Country.Names.En
		meta.CountryCode = city.Country.IsoCode
		meta.Latitude = city.Location.Latitude
		meta.Longitude = city.Location.Longitude
	}

	// Reverse DNS
	resolver := &net.Resolver{}
	names, err := resolver.LookupAddr(ctx, ipStr)
	if err == nil && len(names) > 0 {
		fqdn := strings.TrimSuffix(names[0], ".")
		meta.FQDN = fqdn

		meta.Domain = utils.GetBaseDomain(fqdn)
	}

	return meta, nil
}