iplookup/iplookup.go

package main

import (
	"container/list"
	"context"
	"crypto/tls"
	"flag"
	"fmt"
	"html/template"
	"net"
	"net/http"
	"strings"
	"sync"
	"time"

	"github.com/oschwald/geoip2-golang"
	"golang.org/x/exp/rand"
)

var (
	bind        = flag.String("b", "0.0.0.0", "Bind to address")
	port        = flag.Int("p", 8080, "Port to listen on")
	name        = flag.String("n", "iplookup.fr", "Name of the application")
	domain      = flag.String("d", "iplookup.fr", "Domain name")
	shortdomain = flag.String("s", "iplk.fr", "Short domain name")
	geocity     = flag.String("c", "", "Path to the GeoIP database")
	geoasn      = flag.String("a", "", "Path to the GeoASN database")
	verbose     = flag.Bool("v", false, "Enable verbose logging")
	ttl         = flag.Int("t", 300, "Cache TTL in seconds")
	maxcache    = flag.Int("m", 2048, "Max number of entries in cache")

	// Global database connections (opened once at startup)
	geoCityDB *geoip2.Reader
	geoAsnDB  *geoip2.Reader

	// Cache with map for O(1) lookups and mutex for concurrent access
	cache      = make(map[string]*CacheEntry)
	cacheList  = list.New() // LRU eviction list
	cacheMutex sync.RWMutex

	// Template parsed once at startup
	tmpl *template.Template

	cli = []string{
		"curl",
		"HTTPie",
		"httpie-go",
		"Wget",
		"fetch libfetch",
		"Go",
		"Go-http-client",
		"ddclient",
		"Mikrotik",
		"xh",
	}
	colors = []string{
		"#2488bf",
		"#d84d3d",
		"#f39700",
		"#4caf50",
	}
)

type Ssl struct {
	San       []string
	Ocsp      []string
	Subject   string
	Issuer    string
	Alpn      string
	NotBefore string
	NotAfter  string
}

type Geo struct {
	Country   string
	Ccode     string
	City      string
	Latitude  float64
	Longitude float64
	Postal    string
	Timezone  string
}

type Asn struct {
	Asn uint
	Org string
}

type Info struct {
	Ip               string
	Hostname         string
	Query            bool
	Port             string
	UserAgent        string
	MimeType         string
	Language         string
	ContentType      string
	Encoding         string
	Method           string
	CacheControl     string
	Referer          string
	XForwardedFor    string
	XForwardedProto  string
	Cached           bool
	Ttl              int
	Geo              Geo
	Asn              Asn
	Ssl              Ssl
}

type CacheEntry struct {
	Query   string
	Info    Info
	Time    time.Time
	Element *list.Element // Pointer to element in LRU list for O(1) removal
}

type Tmpl struct {
	Color    string
	Title    string
	Domain   string
	Short    string
	Argument string
	Info     Info
}

func getElem(info Info, elem string) string {
	elem = strings.ToLower(elem)
	elem = strings.ReplaceAll(elem, "-", "_")

	// Use direct field access instead of large switch statement
	switch elem {
	case "hostname", "host":
		return info.Hostname
	case "ip":
		return info.Ip
	case "port":
		return info.Port
	case "user_agent", "ua", "useragent":
		return info.UserAgent
	case "mime_type":
		return info.MimeType
	case "language":
		return info.Language
	case "content_type":
		return info.ContentType
	case "encoding":
		return info.Encoding
	case "method":
		return info.Method
	case "cache_control":
		return info.CacheControl
	case "referer":
		return info.Referer
	case "x_forwarded_for":
		return info.XForwardedFor
	case "x_forwarded_proto":
		return info.XForwardedProto
	case "country":
		return info.Geo.Country
	case "ccode", "cc":
		return info.Geo.Ccode
	case "city":
		return info.Geo.City
	case "latitude":
		return fmt.Sprintf("%f", info.Geo.Latitude)
	case "longitude":
		return fmt.Sprintf("%f", info.Geo.Longitude)
	case "postal":
		return info.Geo.Postal
	case "timezone":
		return info.Geo.Timezone
	case "asn":
		return fmt.Sprintf("%d", info.Asn.Asn)
	case "org", "organisation", "isp":
		return info.Asn.Org
	case "ssl_san", "san":
		return strings.Join(info.Ssl.San, ", ")
	case "ssl_ocsp", "ocsp":
		return strings.Join(info.Ssl.Ocsp, ", ")
	case "ssl_subject", "subject":
		return info.Ssl.Subject
	case "ssl_issuer", "issuer":
		return info.Ssl.Issuer
	case "ssl_alpn", "alpn":
		return info.Ssl.Alpn
	case "ssl_not_after", "not_after":
		return info.Ssl.NotAfter
	case "ssl_not_before", "not_before":
		return info.Ssl.NotBefore
	}
	return ""
}

func writeSsl(w http.ResponseWriter, ssl Ssl) {
	if ssl.San == nil {
		return
	}
	// Direct field access instead of reflection
	if len(ssl.San) > 0 {
		fmt.Fprintf(w, "san: %v\n", ssl.San)
	}
	if len(ssl.Ocsp) > 0 {
		fmt.Fprintf(w, "ocsp: %v\n", ssl.Ocsp)
	}
	if ssl.Subject != "" {
		fmt.Fprintf(w, "subject: %s\n", ssl.Subject)
	}
	if ssl.Issuer != "" {
		fmt.Fprintf(w, "issuer: %s\n", ssl.Issuer)
	}
	if ssl.Alpn != "" {
		fmt.Fprintf(w, "alpn: %s\n", ssl.Alpn)
	}
	if ssl.NotBefore != "" {
		fmt.Fprintf(w, "notbefore: %s\n", ssl.NotBefore)
	}
	if ssl.NotAfter != "" {
		fmt.Fprintf(w, "notafter: %s\n", ssl.NotAfter)
	}
}

func writeAsn(w http.ResponseWriter, asn Asn) {
	if asn.Asn == 0 {
		return
	}
	fmt.Fprintf(w, "asn: %d\n", asn.Asn)
	if asn.Org != "" {
		fmt.Fprintf(w, "org: %s\n", asn.Org)
	}
}

func writeGeo(w http.ResponseWriter, geo Geo) {
	if geo.Country == "" {
		return
	}
	fmt.Fprintf(w, "country: %s\n", geo.Country)
	if geo.Ccode != "" {
		fmt.Fprintf(w, "ccode: %s\n", geo.Ccode)
	}
	if geo.City != "" {
		fmt.Fprintf(w, "city: %s\n", geo.City)
	}
	if geo.Latitude != 0 {
		fmt.Fprintf(w, "latitude: %f\n", geo.Latitude)
	}
	if geo.Longitude != 0 {
		fmt.Fprintf(w, "longitude: %f\n", geo.Longitude)
	}
	if geo.Postal != "" {
		fmt.Fprintf(w, "postal: %s\n", geo.Postal)
	}
	if geo.Timezone != "" {
		fmt.Fprintf(w, "timezone: %s\n", geo.Timezone)
	}
}

func writeInfo(w http.ResponseWriter, info Info) {
	w.Header().Set("Content-Type", "text/plain; charset=utf-8")
	w.WriteHeader(http.StatusOK)
	// Direct field access instead of reflection for better performance
	if info.Ip != "" {
		fmt.Fprintf(w, "ip: %s\n", info.Ip)
	}
	if info.Hostname != "" {
		fmt.Fprintf(w, "hostname: %s\n", info.Hostname)
	}
	if info.Port != "" {
		fmt.Fprintf(w, "port: %s\n", info.Port)
	}
	if info.UserAgent != "" {
		fmt.Fprintf(w, "useragent: %s\n", info.UserAgent)
	}
	if info.MimeType != "" {
		fmt.Fprintf(w, "mimetype: %s\n", info.MimeType)
	}
	if info.Language != "" {
		fmt.Fprintf(w, "language: %s\n", info.Language)
	}
	if info.ContentType != "" {
		fmt.Fprintf(w, "contenttype: %s\n", info.ContentType)
	}
	if info.Encoding != "" {
		fmt.Fprintf(w, "encoding: %s\n", info.Encoding)
	}
	if info.Method != "" {
		fmt.Fprintf(w, "method: %s\n", info.Method)
	}
	if info.CacheControl != "" {
		fmt.Fprintf(w, "cachecontrol: %s\n", info.CacheControl)
	}
	if info.Referer != "" {
		fmt.Fprintf(w, "referer: %s\n", info.Referer)
	}
	if info.XForwardedFor != "" {
		fmt.Fprintf(w, "xforwardedfor: %s\n", info.XForwardedFor)
	}
	if info.XForwardedProto != "" {
		fmt.Fprintf(w, "xforwardedproto: %s\n", info.XForwardedProto)
	}
	if info.Ttl > 0 {
		fmt.Fprintf(w, "ttl: %d\n", info.Ttl)
	}
	writeGeo(w, info.Geo)
	writeAsn(w, info.Asn)
	writeSsl(w, info.Ssl)
}

func reqtype(ip string) string {
	if net.ParseIP(ip) != nil {
		return "ip"
	}
	for i := 0; i < len(ip); i++ {
		if ip[i] == ':' {
			return "ip"
		}
		if ip[i] < '0' || ip[i] > '9' {
			if strings.Contains(ip, ".") {
				return "domain"
			} else {
				return "info"
			}
		}
	}
	return ""
}

func getSslInfo(domain string, info Info) Info {
	lookup, err := net.LookupIP(domain)
	if err != nil {
		if *verbose {
			fmt.Printf("DNS lookup error for %s: %v\n", domain, err)
		}
		return info
	}

	// Check if lookup returned any results
	if len(lookup) == 0 {
		if *verbose {
			fmt.Printf("DNS lookup returned no results for %s\n", domain)
		}
		return info
	}
	info.Ip = lookup[0].String()

	ssl := Ssl{}
	conn, err := tls.DialWithDialer(&net.Dialer{
		Timeout: 2 * time.Second,
	}, "tcp", domain+":443", &tls.Config{InsecureSkipVerify: true})
	if err != nil {
		if *verbose {
			fmt.Printf("TLS connection error for %s: %v\n", domain, err)
		}
		return info
	}
	defer conn.Close()

	// Check if peer certificates are available
	certs := conn.ConnectionState().PeerCertificates
	if len(certs) == 0 {
		if *verbose {
			fmt.Printf("No peer certificates received for %s\n", domain)
		}
		return info
	}

	ssl.San = certs[0].DNSNames
	ssl.Subject = certs[0].Subject.String()
	ssl.Issuer = certs[0].Issuer.String()
	ssl.Alpn = conn.ConnectionState().NegotiatedProtocol
	ssl.Ocsp = certs[0].OCSPServer
	ssl.NotAfter = certs[0].NotAfter.String()
	ssl.NotBefore = certs[0].NotBefore.String()
	info.Ssl = ssl
	return info
}

func getAsn(ip string) Asn {
	asn := Asn{}
	if geoAsnDB == nil {
		return asn
	}

	record, err := geoAsnDB.ASN(net.ParseIP(ip))
	if err != nil {
		if *verbose {
			fmt.Printf("ASN lookup error for %s: %v\n", ip, err)
		}
		return Asn{}
	}

	asn.Asn = record.AutonomousSystemNumber
	asn.Org = record.AutonomousSystemOrganization
	return asn
}

func getGeo(ip string) Geo {
	geo := Geo{}
	if geoCityDB == nil {
		return geo
	}

	record, err := geoCityDB.City(net.ParseIP(ip))
	if err != nil {
		if *verbose {
			fmt.Printf("GeoIP lookup error for %s: %v\n", ip, err)
		}
		return Geo{}
	}

	geo.City = record.City.Names["en"]
	geo.Latitude, geo.Longitude = record.Location.Latitude, record.Location.Longitude
	geo.Country = record.Country.Names["en"]
	geo.Ccode = record.Country.IsoCode
	geo.Postal = record.Postal.Code
	geo.Timezone = record.Location.TimeZone

	return geo
}

func getIP(info Info, r *http.Request) (Info, error) {
	var err error
	if r.Header.Get("X-Forwarded-For") != "" {
		info.Ip = r.Header.Get("X-Forwarded-For")
	} else if r.Header.Get("X-Real-IP") != "" {
		info.Ip = r.Header.Get("X-Real-IP")
	} else if r.Header.Get("X-Client-IP") != "" {
		info.Ip = r.Header.Get("X-Client-IP")
	} else if r.Header.Get("CF-Connecting-IP") != "" {
		info.Ip = r.Header.Get("CF-Connecting-IP")
	} else {
		info.Ip, info.Port, err = net.SplitHostPort(r.RemoteAddr)
		if err != nil {
			if *verbose {
				fmt.Printf("Error parsing remote address %s: %v\n", r.RemoteAddr, err)
			}
			return info, err
		}
	}
	return info, nil
}

func getIpInfo(info Info, r *http.Request) (Info, error) {
	// Create a context with timeout for DNS lookups
	ctx, cancel := context.WithTimeout(r.Context(), 2*time.Second)
	defer cancel()

	// Use a custom resolver with timeout
	resolver := &net.Resolver{}
	hosts, err := resolver.LookupAddr(ctx, info.Ip)
	if err != nil {
		// If lookup fails or times out, use the IP as hostname
		if *verbose {
			fmt.Printf("Reverse DNS lookup error for %s: %v\n", info.Ip, err)
		}
		hosts = []string{info.Ip}
	}

	if len(hosts) > 0 {
		info.Hostname = hosts[0]
	} else {
		info.Hostname = info.Ip
	}

	if !info.Query {
		info.UserAgent = r.UserAgent()
		info.MimeType = r.Header.Get("Accept")
		info.Language = r.Header.Get("Accept-Language")
		info.ContentType = r.Header.Get("Content-Type")
		info.Encoding = r.Header.Get("Accept-Encoding")
		info.Method = r.Method
		info.CacheControl = r.Header.Get("Cache-Control")
		info.Referer = r.Header.Get("Referer")
		info.XForwardedFor = r.Header.Get("X-Forwarded-For")
		info.XForwardedProto = r.Header.Get("X-Forwarded-Proto")
	}
	return info, nil
}

func getCache(query string, info Info) Info {
	cacheMutex.RLock()
	entry, exists := cache[query]
	cacheMutex.RUnlock()

	if !exists {
		return info
	}

	elapsed := time.Since(entry.Time)
	if elapsed < time.Duration(*ttl)*time.Second {
		entry.Info.Ttl = int((time.Duration(*ttl)*time.Second - elapsed).Seconds())
		entry.Info.Cached = true

		// Move to front of LRU list (most recently used)
		cacheMutex.Lock()
		cacheList.MoveToFront(entry.Element)
		cacheMutex.Unlock()

		return entry.Info
	}

	// Entry expired, remove it
	cacheMutex.Lock()
	if entry.Element != nil {
		cacheList.Remove(entry.Element)
	}
	delete(cache, query)
	cacheMutex.Unlock()
	return info
}

func putCache(query string, info Info) {
	info.Cached = true
	info.Ttl = *ttl

	cacheMutex.Lock()
	defer cacheMutex.Unlock()

	// Check if entry already exists (update case)
	if existing, exists := cache[query]; exists {
		existing.Info = info
		existing.Time = time.Now()
		cacheList.MoveToFront(existing.Element)
		return
	}

	// LRU eviction if cache is full
	if len(cache) >= *maxcache {
		// Remove least recently used (back of list)
		oldest := cacheList.Back()
		if oldest != nil {
			oldEntry := oldest.Value.(*CacheEntry)
			delete(cache, oldEntry.Query)
			cacheList.Remove(oldest)
		}
	}

	// Add new entry
	entry := &CacheEntry{
		Query: query,
		Info:  info,
		Time:  time.Now(),
	}
	entry.Element = cacheList.PushFront(entry)
	cache[query] = entry
}

func handler(w http.ResponseWriter, r *http.Request) {
	info := Info{}
	arg := ""
	isCli := false
	var err error

	if *verbose {
		fmt.Println("Request:", r)
	}

	if r.Method != "GET" {
		http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
		return
	}

	for i := 0; i < len(cli); i++ {
		if strings.Contains(r.UserAgent(), cli[i]) {
			isCli = true
		}
	}

	if r.URL.Path == "/" {
		info, err = getIP(info, r)
		if err != nil {
			http.Error(w, err.Error(), http.StatusInternalServerError)
			if *verbose {
				fmt.Printf("Error getting IP: %v\n", err)
			}
			return
		}
		if isCli {
			w.Header().Set("Content-Type", "text/plain; charset=utf-8")
			w.WriteHeader(http.StatusOK)
			w.Write([]byte(info.Ip + "\n"))
			return
		}
		info = getCache(info.Ip, info)
		w.Header().Set("Cache-Control", fmt.Sprintf("max-age=%d", info.Ttl))
		if !info.Cached {
			if *geocity != "" {
				info.Geo = getGeo(info.Ip)
			}
			if *geoasn != "" {
				info.Asn = getAsn(info.Ip)
			}
			putCache(info.Ip, info)
			w.Header().Set("Cache-Control", fmt.Sprintf("max-age=%d", *ttl))
		}
		info, err = getIpInfo(info, r)
		if err != nil {
			http.Error(w, err.Error(), http.StatusInternalServerError)
			if *verbose {
				fmt.Printf("Error getting IP info: %v\n", err)
			}
			return
		}
	}

	if r.URL.Path == "/favicon.ico" {
		w.Header().Set("Content-Type", "image/png")
		http.ServeFile(w, r, "favicon.png")
		return
	}

	args := strings.Split(r.URL.Path, "/")
	if len(args) > 1 && args[1] != "" {
		// Validate that args[1] is not excessively long (prevent DoS)
		if len(args[1]) > 256 {
			http.Error(w, "Request URI too long", http.StatusRequestURITooLong)
			return
		}

		arg = "/" + args[1]
		rtype := reqtype(args[1])
		if rtype == "" {
			http.Error(w, "Not found", http.StatusNotFound)
			return
		}
		if rtype == "ip" {
			info.Query = true
			info.Ip = args[1]
			info = getCache(args[1], info)
		} else if rtype == "domain" {
			info.Query = true
			info = getCache(args[1], info)
		} else {
			info, err = getIP(info, r)
			if err != nil {
				http.Error(w, err.Error(), http.StatusInternalServerError)
				if *verbose {
					fmt.Printf("Error getting IP: %v\n", err)
				}
				return
			}
			info = getCache(info.Ip, info)
		}

		if !info.Cached {
			if rtype == "domain" {
				info.Query = true
				info = getSslInfo(args[1], info)
			}
			if info.Ip == "" {
				http.Error(w, "Not found", http.StatusNotFound)
				return
			}
			if *geocity != "" {
				info.Geo = getGeo(info.Ip)
			}
			if *geoasn != "" {
				info.Asn = getAsn(info.Ip)
			}

			if rtype != "info" {
				putCache(args[1], info)
			} else {
				putCache(info.Ip, info)
			}
		}

		w.Header().Set("Cache-Control", fmt.Sprintf("max-age=%d", info.Ttl))

		info, err = getIpInfo(info, r)
		if err != nil {
			http.Error(w, err.Error(), http.StatusInternalServerError)
			if *verbose {
				fmt.Printf("Error getting IP info: %v\n", err)
			}
			return
		}

		if rtype == "info" || (len(args) > 2 && args[2] != "") {
			// Safely access the last argument
			lastArg := args[len(args)-1]
			if lastArg == "all" {
				writeInfo(w, info)
				return
			}
			if lastArg == "geo" {
				writeGeo(w, info.Geo)
				return
			}
			if lastArg == "as" {
				writeAsn(w, info.Asn)
				return
			}
			if lastArg == "ssl" {
				writeSsl(w, info.Ssl)
				return
			}
			res := getElem(info, lastArg)
			if res == "" {
				http.Error(w, "Not found", http.StatusNotFound)
				return
			}
			w.Header().Set("Content-Type", "text/plain; charset=utf-8")
			w.WriteHeader(http.StatusOK)
			fmt.Fprintf(w, "%s\n", res)
			return
		}
	}

	if isCli {
		writeInfo(w, info)
		return
	}

	tmplData := Tmpl{
		Title:    *name,
		Domain:   *domain,
		Short:    *shortdomain,
		Color:    colors[rand.Intn(len(colors))],
		Argument: arg,
		Info:     info,
	}

	if tmpl == nil {
		http.Error(w, "Template not initialized", http.StatusInternalServerError)
		return
	}

	// Set Content-Type header for HTML response
	w.Header().Set("Content-Type", "text/html; charset=utf-8")
	err = tmpl.Execute(w, tmplData)
	if err != nil {
		if *verbose {
			fmt.Printf("Template execution error: %v\n", err)
		}
		return
	}
}

func initialize() error {
	var err error

	// Open GeoIP City database if path is provided
	if *geocity != "" {
		geoCityDB, err = geoip2.Open(*geocity)
		if err != nil {
			return fmt.Errorf("failed to open GeoIP City database: %w", err)
		}
		fmt.Println("GeoIP City database loaded:", *geocity)
	}

	// Open GeoIP ASN database if path is provided
	if *geoasn != "" {
		geoAsnDB, err = geoip2.Open(*geoasn)
		if err != nil {
			return fmt.Errorf("failed to open GeoIP ASN database: %w", err)
		}
		fmt.Println("GeoIP ASN database loaded:", *geoasn)
	}

	// Parse template once at startup
	tmpl, err = template.ParseFiles("template.html")
	if err != nil {
		return fmt.Errorf("failed to parse template: %w", err)
	}
	fmt.Println("Template loaded: template.html")

	return nil
}

func main() {
	flag.Parse()

	// Initialize databases and template
	if err := initialize(); err != nil {
		fmt.Println("Initialization error:", err)
		return
	}
	defer func() {
		if geoCityDB != nil {
			geoCityDB.Close()
		}
		if geoAsnDB != nil {
			geoAsnDB.Close()
		}
	}()

	fmt.Println("Listening on port", *port)
	http.HandleFunc("/", handler)
	err := http.ListenAndServe(fmt.Sprintf("%s:%d", *bind, *port), nil)
	if err != nil {
		fmt.Println(err)
	}
}