IPDATA/ipstore.go

package main

import (
	"context"
	"encoding/csv"
	"errors"
	"fmt"
	"io"
	"log"
	"math/big"
	"net"
	"os"
	"sort"
	"strconv"
	"strings"
	"time"
)

// ASNInfo represents data from the ASN CSV file
type ASNInfo struct {
	IPRangeStart *big.Int
	IPRangeEnd   *big.Int
	CIDR         string
	ASN          string
	ASName       string
}

// GeoInfo represents data from the geolocation CSV file
type GeoInfo struct {
	IPRangeStart *big.Int
	IPRangeEnd   *big.Int
	CountryCode  string
	Country      string
	Region       string
	City         string
	Latitude     float64
	Longitude    float64
	ZipCode      string
	TimeZone     string
}

// IPInfo combines ASN and geolocation information
type IPInfo struct {
	Source      string  `json:"source"`
	IP          string  `json:"ip"`
	Hostname    string  `json:"hostname,omitempty"`
	ASN         string  `json:"asn,omitempty"`
	ASName      string  `json:"as_name,omitempty"`
	CountryCode string  `json:"country_code,omitempty"`
	Country     string  `json:"country,omitempty"`
	Region      string  `json:"region,omitempty"`
	City        string  `json:"city,omitempty"`
	Latitude    float64 `json:"latitude,omitempty"`
	Longitude   float64 `json:"longitude,omitempty"`
	ZipCode     string  `json:"zip_code,omitempty"`
	TimeZone    string  `json:"time_zone,omitempty"`
}

type IPsInfo struct {
	Items []IPInfo `json:"items"`
}

// IPStore holds the IP data in memory
type IPStore struct {
	asnDataIPv4  []ASNInfo
	asnDataIPv6  []ASNInfo
	geoDataIPv4  []GeoInfo
	geoDataIPv6  []GeoInfo
}

// NewIPStore creates and initializes a new IPStore
func NewIPStore(asnFileIPv6, geoFileIPv6, asnFileIPv4, geoFileIPv4 string) (*IPStore, error) {
	store := &IPStore{
		asnDataIPv4:  []ASNInfo{},
		asnDataIPv6:  []ASNInfo{},
		geoDataIPv4:  []GeoInfo{},
		geoDataIPv6:  []GeoInfo{},
	}

	// Load IPv6 ASN data
	if err := store.loadASNData(asnFileIPv6, &store.asnDataIPv6); err != nil {
		return nil, fmt.Errorf("failed to load IPv6 ASN data: %w", err)
	}

	// Load IPv6 geolocation data
	if err := store.loadGeoData(geoFileIPv6, &store.geoDataIPv6); err != nil {
		return nil, fmt.Errorf("failed to load IPv6 geolocation data: %w", err)
	}

	// Load IPv4 ASN data
	if err := store.loadASNData(asnFileIPv4, &store.asnDataIPv4); err != nil {
		return nil, fmt.Errorf("failed to load IPv4 ASN data: %w", err)
	}

	// Load IPv4 geolocation data
	if err := store.loadGeoData(geoFileIPv4, &store.geoDataIPv4); err != nil {
		return nil, fmt.Errorf("failed to load IPv4 geolocation data: %w", err)
	}

	// Sort data for binary search (IPv6)
	sort.Slice(store.asnDataIPv6, func(i, j int) bool {
		return store.asnDataIPv6[i].IPRangeStart.Cmp(store.asnDataIPv6[j].IPRangeStart) < 0
	})

	sort.Slice(store.geoDataIPv6, func(i, j int) bool {
		return store.geoDataIPv6[i].IPRangeStart.Cmp(store.geoDataIPv6[j].IPRangeStart) < 0
	})

	// Sort data for binary search (IPv4)
	sort.Slice(store.asnDataIPv4, func(i, j int) bool {
		return store.asnDataIPv4[i].IPRangeStart.Cmp(store.asnDataIPv4[j].IPRangeStart) < 0
	})

	sort.Slice(store.geoDataIPv4, func(i, j int) bool {
		return store.geoDataIPv4[i].IPRangeStart.Cmp(store.geoDataIPv4[j].IPRangeStart) < 0
	})

	return store, nil
}

// loadASNData parses the ASN CSV file
func (s *IPStore) loadASNData(filename string, data *[]ASNInfo) error {
	file, err := os.Open(filename)
	if err != nil {
		return err
	}
	defer file.Close()

	reader := csv.NewReader(file)
	reader.Comma = ','

	for {
		record, err := reader.Read()
		if err == io.EOF {
			break
		}
		if err != nil {
			return err
		}

		if len(record) < 5 {
			continue // Skip malformed records
		}

		// Parse IP range
		ipStart := new(big.Int)
		ipStart.SetString(record[0], 10)

		ipEnd := new(big.Int)
		ipEnd.SetString(record[1], 10)

		*data = append(*data, ASNInfo{
			IPRangeStart: ipStart,
			IPRangeEnd:   ipEnd,
			CIDR:         record[2],
			ASN:          record[3],
			ASName:       record[4],
		})
	}

	return nil
}

// loadGeoData parses the geolocation CSV file
func (s *IPStore) loadGeoData(filename string, data *[]GeoInfo) error {
	file, err := os.Open(filename)
	if err != nil {
		return err
	}
	defer file.Close()

	reader := csv.NewReader(file)
	reader.Comma = ','

	for {
		record, err := reader.Read()
		if err == io.EOF {
			break
		}
		if err != nil {
			return err
		}

		if len(record) < 10 {
			continue // Skip malformed records
		}

		// Parse IP range
		ipStart := new(big.Int)
		ipStart.SetString(record[0], 10)

		ipEnd := new(big.Int)
		ipEnd.SetString(record[1], 10)

		// Parse coordinates
		lat, _ := strconv.ParseFloat(record[6], 64)
		long, _ := strconv.ParseFloat(record[7], 64)

		*data = append(*data, GeoInfo{
			IPRangeStart: ipStart,
			IPRangeEnd:   ipEnd,
			CountryCode:  record[2],
			Country:      record[3],
			Region:       record[4],
			City:         record[5],
			Latitude:     lat,
			Longitude:    long,
			ZipCode:      record[8],
			TimeZone:     record[9],
		})
	}

	return nil
}

// lookupHostname performs a reverse DNS lookup to find the hostname for an IP
// with a timeout to avoid long waiting times
func lookupHostname(ipStr string, timeout int) string {
	// Check if DNS lookup is disabled
	if disableDNS := os.Getenv("DISABLE_DNS_LOOKUP"); disableDNS != "" {
		if strings.ToLower(disableDNS) == "true" || disableDNS == "1" {
			return ""
		}
	}
	
	startTime := time.Now()
	
	// Get timeout from environment or use default
	timeoutMs := timeout
	if envTimeout := os.Getenv("DNS_TIMEOUT_MS"); envTimeout != "" {
		if t, err := strconv.Atoi(envTimeout); err == nil && t > 0 {
			timeoutMs = t
		}
	}

	// Create a context with the configured timeout
	ctx, cancel := context.WithTimeout(context.Background(), time.Duration(timeoutMs)*time.Millisecond)
	defer cancel()

	// Use a channel to handle the timeout in parallel with the DNS lookup
	resultChan := make(chan string, 1)
	
	go func() {
		// Perform the DNS lookup in a goroutine
		names, err := net.DefaultResolver.LookupAddr(ctx, ipStr)
		if err != nil {
			if ctx.Err() == context.DeadlineExceeded {
				log.Printf("DNS lookup timeout for IP %s (limit: %dms)", ipStr, timeoutMs)
			} else {
				log.Printf("DNS lookup error for IP %s: %v", ipStr, err)
			}
			resultChan <- ""
			return
		}
		if len(names) == 0 {
			resultChan <- ""
			return
		}
		resultChan <- names[0]
	}()

	// Wait for the result or timeout
	select {
	case result := <-resultChan:
		elapsed := time.Since(startTime)
		if result != "" {
			log.Printf("DNS lookup for IP %s took %dms: found hostname %s", 
				ipStr, elapsed.Milliseconds(), result)
		}
		return result
	case <-ctx.Done():
		// Timeout or context cancelled
		log.Printf("DNS lookup for IP %s timed out after %dms", 
			ipStr, time.Since(startTime).Milliseconds())
		return ""
	}
}

// ipToBigInt converts an IP address to a big.Int
func ipToBigInt(ipStr string) (*big.Int, bool, error) {
	ip := net.ParseIP(ipStr)
	if ip == nil {
		return nil, false, errors.New("invalid IP address")
	}

	// Check if it's IPv4 or IPv6
	isIPv4 := ip.To4() != nil

	// For IPv4, ensure it's in IPv4 format (not mapped)
	if isIPv4 {
		ip = ip.To4()
	}

	ipInt := new(big.Int)
	ipBytes := []byte(ip)
	ipInt.SetBytes(ipBytes)

	return ipInt, isIPv4, nil
}

// findASNInfo uses binary search to find ASN info for an IP
func (s *IPStore) findASNInfo(ipInt *big.Int, isIPv4 bool) *ASNInfo {
	var data []ASNInfo
	if isIPv4 {
		data = s.asnDataIPv4
	} else {
		data = s.asnDataIPv6
	}

	idx := sort.Search(len(data), func(i int) bool {
		return data[i].IPRangeEnd.Cmp(ipInt) >= 0
	})

	if idx < len(data) && 
	   data[idx].IPRangeStart.Cmp(ipInt) <= 0 && 
	   data[idx].IPRangeEnd.Cmp(ipInt) >= 0 {
		return &data[idx]
	}

	return nil
}

// findGeoInfo uses binary search to find geolocation info for an IP
func (s *IPStore) findGeoInfo(ipInt *big.Int, isIPv4 bool) *GeoInfo {
	var data []GeoInfo
	if isIPv4 {
		data = s.geoDataIPv4
	} else {
		data = s.geoDataIPv6
	}

	idx := sort.Search(len(data), func(i int) bool {
		return data[i].IPRangeEnd.Cmp(ipInt) >= 0
	})

	if idx < len(data) && 
	   data[idx].IPRangeStart.Cmp(ipInt) <= 0 && 
	   data[idx].IPRangeEnd.Cmp(ipInt) >= 0 {
		return &data[idx]
	}

	return nil
}

// GetIPInfo retrieves information about an IP address
func (s *IPStore) GetIPInfo(ipStr string, timeout int, source string) (*IPInfo, error) {
	ipInt, isIPv4, err := ipToBigInt(ipStr)
	if err != nil {
		return nil, err
	}

	info := &IPInfo{IP: ipStr, Source: source}
	
	// Add hostname via reverse DNS lookup
	info.Hostname = lookupHostname(ipStr, timeout)

	// Find ASN info using binary search
	if asnInfo := s.findASNInfo(ipInt, isIPv4); asnInfo != nil {
		info.ASN = asnInfo.ASN
		info.ASName = asnInfo.ASName
	}

	// Find geolocation info using binary search
	if geoInfo := s.findGeoInfo(ipInt, isIPv4); geoInfo != nil {
		info.CountryCode = geoInfo.CountryCode
		info.Country = geoInfo.Country
		info.Region = geoInfo.Region
		info.City = geoInfo.City
		info.Latitude = geoInfo.Latitude
		info.Longitude = geoInfo.Longitude
		info.ZipCode = geoInfo.ZipCode
		info.TimeZone = geoInfo.TimeZone
	}

	return info, nil
}

// IsValidIP checks if the provided string is a valid IP address
func (s *IPStore) IsValidIP(input string) bool {
	return net.ParseIP(input) != nil
}

// ResolveDomain resolves a domain name to an IP address
func (s *IPStore) ResolveDomain(domain string, timeout int) (string, error) {
	if envTimeout := os.Getenv("DNS_TIMEOUT_MS"); envTimeout != "" {
		if t, err := strconv.Atoi(envTimeout); err == nil && t > 0 {
			timeout = t
		}
	}

	ctx, cancel := context.WithTimeout(context.Background(), time.Duration(timeout)*time.Millisecond)
	defer cancel()

	// Use net.DefaultResolver to perform the lookup
	ips, err := net.DefaultResolver.LookupIP(ctx, "ip", domain)
	if err != nil {
		return "", fmt.Errorf("failed to resolve domain %s: %w", domain, err)
	}

	if len(ips) == 0 {
		return "", fmt.Errorf("no IP addresses found for domain %s", domain)
	}

	// Prefer IPv4 addresses if available
	for _, ip := range ips {
		if ip.To4() != nil {
			return ip.String(), nil
		}
	}

	// Return the first IP address if no IPv4 available
	return ips[0].String(), nil
}