subconverter-go/internal/conversion/geoip.go

package conversion

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"net"
	"net/http"
	"net/url"
	"path"
	"strings"
	"sync"
	"time"

	"github.com/subconverter-go/internal/logging"
)

// GeoIPInfo represents a subset of geoip info exposed to scripts.
type GeoIPInfo struct {
	IP          string  `json:"ip"`
	CountryCode string  `json:"country_code"`
	CountryName string  `json:"country_name"`
	Region      string  `json:"region"`
	City        string  `json:"city"`
	ISP         string  `json:"isp"`
	ASN         string  `json:"asn"`
	Latitude    float64 `json:"latitude,omitempty"`
	Longitude   float64 `json:"longitude,omitempty"`
}

// GeoIPResolver resolves geo-location information for a given address.
type GeoIPResolver interface {
	Lookup(ctx context.Context, address string) (*GeoIPInfo, error)
}

const (
	defaultGeoIPBaseURL       = "https://api.ip.sb/geoip"
	geoIPRequestTimeout       = 3 * time.Second
	geoIPUserAgent            = "subconverter-go/geoip"
	classificationPrivate     = "PRIVATE"
	classificationUnspecified = "UNSPECIFIED"
)

type httpGeoIPResolver struct {
	baseURL string
	client  *http.Client
	logger  *logging.Logger
}

func newDefaultGeoIPResolver(logger *logging.Logger) GeoIPResolver {
	return &httpGeoIPResolver{
		baseURL: defaultGeoIPBaseURL,
		client: &http.Client{
			Timeout: geoIPRequestTimeout,
		},
		logger: logger,
	}
}

// SetBaseURL allows overriding the base geoip endpoint (used for tests).
func (r *httpGeoIPResolver) SetBaseURL(base string) {
	if base != "" {
		r.baseURL = base
	}
}

func (r *httpGeoIPResolver) Lookup(ctx context.Context, address string) (*GeoIPInfo, error) {
	target, err := normalizeGeoIPTarget(address)
	if err != nil {
		return nil, err
	}

	if info := classifySpecialIP(target); info != nil {
		return info, nil
	}

	if r == nil {
		return nil, errors.New("geoip resolver is not configured")
	}

	endpoint, err := buildGeoIPURL(r.baseURL, target)
	if err != nil {
		return nil, err
	}

	req, err := http.NewRequestWithContext(ctx, http.MethodGet, endpoint, nil)
	if err != nil {
		return nil, err
	}
	req.Header.Set("User-Agent", geoIPUserAgent)

	resp, err := r.client.Do(req)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		return &GeoIPInfo{IP: target}, fmt.Errorf("geoip lookup failed with status %d", resp.StatusCode)
	}

	var payload GeoIPInfo
	decoder := json.NewDecoder(resp.Body)
	if err := decoder.Decode(&payload); err != nil {
		return nil, err
	}

	if payload.IP == "" {
		payload.IP = target
	}
	return &payload, nil
}

func normalizeGeoIPTarget(input string) (string, error) {
	trimmed := strings.TrimSpace(input)
	if trimmed == "" {
		return "", errors.New("empty address")
	}

	// If input is a URL, extract host part.
	if strings.Contains(trimmed, "://") {
		if parsed, err := url.Parse(trimmed); err == nil && parsed.Host != "" {
			trimmed = parsed.Host
		}
	}

	// Remove port if present.
	if host, _, err := net.SplitHostPort(trimmed); err == nil && host != "" {
		trimmed = host
	}

	trimmed = strings.Trim(trimmed, "[]")
	return trimmed, nil
}

func buildGeoIPURL(baseURL, target string) (string, error) {
	if baseURL == "" {
		return "", errors.New("geoip base url is empty")
	}

	parsed, err := url.Parse(baseURL)
	if err != nil {
		return "", err
	}

	// Ensure trailing slash and append escaped target.
	parsed.Path = path.Join(parsed.Path, url.PathEscape(target))
	return parsed.String(), nil
}

func classifySpecialIP(target string) *GeoIPInfo {
	ip := net.ParseIP(target)
	if ip == nil {
		return nil
	}

	switch {
	case ip.IsLoopback(), ip.IsUnspecified():
		return &GeoIPInfo{
			IP:          ip.String(),
			CountryCode: classificationUnspecified,
			CountryName: "Unspecified",
			ISP:         "Localhost",
		}
	case ip.IsPrivate(), ip.IsLinkLocalMulticast(), ip.IsLinkLocalUnicast():
		return &GeoIPInfo{
			IP:          ip.String(),
			CountryCode: classificationPrivate,
			CountryName: "Private Network",
			ISP:         "Local Network",
		}
	case ip.IsMulticast():
		return &GeoIPInfo{
			IP:          ip.String(),
			CountryCode: classificationUnspecified,
			CountryName: "Multicast",
			ISP:         "Multicast",
		}
	default:
		return nil
	}
}

func (ce *ConversionEngine) SetGeoIPResolver(resolver GeoIPResolver) {
	if resolver == nil {
		ce.geoResolver = newDefaultGeoIPResolver(ce.logger)
	} else {
		ce.geoResolver = resolver
	}
	ce.geoipCache = sync.Map{}
}

func (ce *ConversionEngine) lookupGeoIP(address string) (*GeoIPInfo, error) {
	target, err := normalizeGeoIPTarget(address)
	if err != nil || target == "" {
		return &GeoIPInfo{IP: target}, err
	}

	if value, ok := ce.geoipCache.Load(target); ok {
		if info, ok := value.(*GeoIPInfo); ok {
			return info, nil
		}
	}

	resolver := ce.geoResolver
	if resolver == nil {
		resolver = newDefaultGeoIPResolver(ce.logger)
		ce.geoResolver = resolver
	}

	ctx, cancel := context.WithTimeout(context.Background(), geoIPRequestTimeout)
	defer cancel()

	info, err := resolver.Lookup(ctx, target)
	if err != nil {
		ce.logger.WithError(err).Debug("geoip lookup failed")
		if classified := classifySpecialIP(target); classified != nil {
			info = classified
		} else if info == nil {
			info = &GeoIPInfo{IP: target}
		}
	}

	if info != nil {
		ce.geoipCache.Store(target, info)
		return info, nil
	}

	return &GeoIPInfo{IP: target}, err
}

func (ce *ConversionEngine) geoIPJSON(address string) string {
	info, err := ce.lookupGeoIP(address)
	if err != nil {
		ce.logger.WithError(err).Debug("geoip lookup returned error")
	}
	data, err := json.Marshal(info)
	if err != nil {
		ce.logger.WithError(err).Warn("failed to marshal geoip info")
		return "{}"
	}
	return string(data)
}