gists/projects/go-tools/goipgrep/cmd/ipgrep/main.go

package main

import (
	"context"
	"encoding/csv"
	"encoding/json"
	"errors"
	"flag"
	"fmt"
	"os"
	"os/signal"
	"runtime"
	"sort"
	"strings"
	"syscall"
	"time"

	"git.ktf.ninja/tabledevil/gists/projects/go-tools/go/goipgrep/internal/cache"
	"git.ktf.ninja/tabledevil/gists/projects/go-tools/go/goipgrep/internal/extract"
	"git.ktf.ninja/tabledevil/gists/projects/go-tools/go/goipgrep/internal/ipinfo"
	"git.ktf.ninja/tabledevil/gists/projects/go-tools/go/goipgrep/internal/normalize"
	"git.ktf.ninja/tabledevil/gists/projects/go-tools/go/goipgrep/internal/ping"
	"git.ktf.ninja/tabledevil/gists/projects/go-tools/go/goipgrep/internal/resolve"
)

var (
	version = "dev"
	commit  = "unknown"
	date    = "unknown"
)

type stringSlice []string

func (s *stringSlice) String() string { return strings.Join(*s, ",") }
func (s *stringSlice) Set(v string) error {
	*s = append(*s, v)
	return nil
}

func main() {
	var (
		// legacy short flags (kept for compatibility)
		shortUniq    bool
		shortSort    bool
		shortMAC     bool
		shortPing    bool
		shortResolve bool
		shortLookup  bool
		shortFile    string

		// long flags
		longUniq    bool
		longSort    bool
		longMAC     bool
		longPing    bool
		longResolve bool
		longLookup  bool
		longFiles   stringSlice

		format      string
		jobs        int
		timeout     time.Duration
		pingMode    string
		ipv6        bool
		cachePath   string
		noCache     bool
		clearCache  bool
		cacheTTL    time.Duration
		cacheMax    int
		ipinfoToken string
		showVersion bool
	)

	flag.BoolVar(&shortUniq, "u", false, "only show uniq IPs/MACs (implies -s)")
	flag.BoolVar(&shortSort, "s", false, "sort output")
	flag.BoolVar(&shortMAC, "m", false, "grep MAC addresses instead of IP addresses")
	flag.BoolVar(&shortPing, "p", false, "only show entries considered reachable (see --ping-mode)")
	flag.BoolVar(&shortResolve, "r", false, "resolve: IP reverse DNS; MAC best-effort via neighbor/ARP table (Linux only)")
	flag.BoolVar(&shortLookup, "l", false, "lookup ip info using ipinfo.io (IP only)")
	flag.StringVar(&shortFile, "f", "", "input file (legacy; prefer --file)")

	flag.BoolVar(&longUniq, "uniq", false, "only show uniq IPs/MACs (implies --sort)")
	flag.BoolVar(&longSort, "sort", false, "sort output")
	flag.BoolVar(&longMAC, "mac", false, "grep MAC addresses instead of IP addresses")
	flag.BoolVar(&longPing, "pingable", false, "only show entries considered reachable (see --ping-mode)")
	flag.BoolVar(&longResolve, "resolve", false, "resolve: IP reverse DNS; MAC best-effort via neighbor/ARP table (Linux only)")
	flag.BoolVar(&longLookup, "lookup", false, "lookup ip info using ipinfo.io (IP only)")
	flag.Var(&longFiles, "file", "input file (repeatable)")

	flag.StringVar(&format, "format", "text", "output format: text|csv|json")
	flag.IntVar(&jobs, "jobs", runtime.GOMAXPROCS(0), "max concurrent workers for network operations")
	flag.DurationVar(&timeout, "timeout", 2*time.Second, "per-operation timeout for -p/-r/-l")
	flag.StringVar(&pingMode, "ping-mode", "auto", "reachability mode for -p: auto|icmp|tcp")
	flag.BoolVar(&ipv6, "ipv6", false, "also match IPv6 addresses (default: IPv4 only)")

	flag.StringVar(&cachePath, "cache-path", "", "cache path for ipinfo lookups (default: OS cache dir)")
	flag.BoolVar(&noCache, "no-cache", false, "disable reading/writing the ipinfo cache")
	flag.BoolVar(&clearCache, "clear-cache", false, "clear the ipinfo cache then exit")
	flag.DurationVar(&cacheTTL, "cache-ttl", 30*24*time.Hour, "ipinfo cache TTL (0 disables expiry)")
	flag.IntVar(&cacheMax, "cache-max-entries", 50000, "max cached IPs for ipinfo lookups (0 disables pruning)")
	flag.StringVar(&ipinfoToken, "ipinfo-token", "", "ipinfo.io token (also read from IPINFO_TOKEN)")

	flag.BoolVar(&showVersion, "version", false, "print version and exit")

	flag.Usage = func() {
		fmt.Fprintf(os.Stderr, "Usage: %s [flags] [file...]\n\n", os.Args[0])
		flag.PrintDefaults()
	}
	flag.Parse()

	if showVersion {
		fmt.Printf("ipgrep %s (%s) %s\n", version, commit, date)
		return
	}

	uniqFlag := shortUniq || longUniq
	sortFlag := shortSort || longSort
	macFlag := shortMAC || longMAC
	pingable := shortPing || longPing
	resolveFlag := shortResolve || longResolve
	lookupFlag := shortLookup || longLookup

	// keep old behavior: pingable/lookup implies uniq+sort
	if pingable || lookupFlag {
		sortFlag = true
		uniqFlag = true
	}
	if uniqFlag {
		sortFlag = true
	}

	if lookupFlag && macFlag {
		fatalf("lookup mode (-l/--lookup) only works for IP addresses, not MAC addresses")
	}

	format = strings.ToLower(strings.TrimSpace(format))
	switch format {
	case "text", "csv", "json":
	default:
		fatalf("invalid --format %q (want text|csv|json)", format)
	}
	pingMode = strings.ToLower(strings.TrimSpace(pingMode))
	switch pingMode {
	case "auto", "icmp", "tcp":
	default:
		fatalf("invalid --ping-mode %q (want auto|icmp|tcp)", pingMode)
	}
	if jobs < 1 {
		fatalf("--jobs must be >= 1")
	}
	if timeout <= 0 {
		fatalf("--timeout must be > 0")
	}

	var files []string
	if shortFile != "" {
		files = append(files, shortFile)
	}
	files = append(files, longFiles...)
	files = append(files, flag.Args()...)

	ctx, cancel := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM)
	defer cancel()

	// cache management for lookup mode
	if lookupFlag && ipinfoToken == "" {
		ipinfoToken = os.Getenv("IPINFO_TOKEN")
	}

	// Resolve cache default path.
	if cachePath == "" {
		cachePath = cache.DefaultPath()
	}

	if clearCache {
		if err := cache.Clear(cachePath); err != nil && !errors.Is(err, os.ErrNotExist) {
			fatalf("clear cache: %v", err)
		}
		return
	}

	mode := extract.ModeIP
	if macFlag {
		mode = extract.ModeMAC
	}

	// Extraction stage (stream input, collect matches as needed).
	var matches []string
	seen := map[string]struct{}{}
	emit := func(s string) {
		if uniqFlag {
			if _, ok := seen[s]; ok {
				return
			}
			seen[s] = struct{}{}
		}
		matches = append(matches, s)
	}

	exOpts := extract.Options{Mode: mode, IPv6: ipv6}

	if len(files) == 0 {
		if err := extract.Grep(ctx, os.Stdin, exOpts, func(raw string) {
			if mode == extract.ModeMAC {
				if norm, ok := normalize.NormalizeMAC(raw); ok {
					emit(norm)
				}
				return
			}
			if norm, ok := normalize.NormalizeIP(raw, ipv6); ok {
				emit(norm)
			}
		}); err != nil {
			fatalf("read stdin: %v", err)
		}
	} else {
		for _, f := range files {
			r, err := os.Open(f)
			if err != nil {
				fmt.Fprintf(os.Stderr, "ipgrep: open %s: %v\n", f, err)
				continue
			}
			err = extract.Grep(ctx, r, exOpts, func(raw string) {
				if mode == extract.ModeMAC {
					if norm, ok := normalize.NormalizeMAC(raw); ok {
						emit(norm)
					}
					return
				}
				if norm, ok := normalize.NormalizeIP(raw, ipv6); ok {
					emit(norm)
				}
			})
			_ = r.Close()
			if err != nil {
				fmt.Fprintf(os.Stderr, "ipgrep: read %s: %v\n", f, err)
			}
		}
	}

	if len(matches) == 0 {
		return
	}

	if sortFlag {
		sort.Strings(matches)
	}

	// -p reachability filter
	if pingable {
		popts := ping.Options{
			Mode:    pingMode,
			Timeout: timeout,
			Jobs:    jobs,
		}
		if macFlag {
			var err error
			matches, err = ping.FilterMACs(ctx, matches, popts)
			if err != nil {
				fatalf("%v", err)
			}
		} else {
			var err error
			matches, err = ping.FilterIPs(ctx, matches, popts)
			if err != nil {
				fatalf("%v", err)
			}
		}
		if sortFlag {
			sort.Strings(matches)
		}
	}

	// -l ipinfo lookup
	if lookupFlag {
		if format == "csv" {
			w := csv.NewWriter(os.Stdout)
			must(w.Write([]string{"ip", "country", "region", "city", "org", "hostname"}))
			client := ipinfo.Client{
				Token:   ipinfoToken,
				Timeout: timeout,
			}
			var store *cache.Store
			if !noCache {
				s, err := cache.Load(cachePath, cacheTTL, cacheMax)
				if err != nil {
					fmt.Fprintf(os.Stderr, "ipgrep: cache load: %v\n", err)
				} else {
					store = s
				}
			}
			results := ipinfo.LookupAll(ctx, matches, client, store, jobs)
			sort.Slice(results, func(i, j int) bool { return results[i].IP < results[j].IP })
			for _, info := range results {
				must(w.Write([]string{info.IP, info.Country, info.Region, info.City, info.Org, info.Hostname}))
			}
			w.Flush()
			if err := w.Error(); err != nil {
				fatalf("csv: %v", err)
			}
			if store != nil && store.Changed() {
				if err := store.Save(); err != nil {
					fmt.Fprintf(os.Stderr, "ipgrep: cache save: %v\n", err)
				}
			}
			return
		}

		client := ipinfo.Client{
			Token:   ipinfoToken,
			Timeout: timeout,
		}
		var store *cache.Store
		if !noCache {
			s, err := cache.Load(cachePath, cacheTTL, cacheMax)
			if err != nil {
				fmt.Fprintf(os.Stderr, "ipgrep: cache load: %v\n", err)
			} else {
				store = s
			}
		}
		results := ipinfo.LookupAll(ctx, matches, client, store, jobs)
		sort.Slice(results, func(i, j int) bool { return results[i].IP < results[j].IP })

		if format == "json" {
			enc := json.NewEncoder(os.Stdout)
			enc.SetIndent("", "  ")
			must(enc.Encode(results))
		} else {
			for _, info := range results {
				fmt.Printf("%s\t%s\t%s\t%s\t%s\t%s\n", info.IP, info.Country, info.Region, info.City, info.Org, info.Hostname)
			}
		}
		if store != nil && store.Changed() {
			if err := store.Save(); err != nil {
				fmt.Fprintf(os.Stderr, "ipgrep: cache save: %v\n", err)
			}
		}
		return
	}

	// -r resolution
	if resolveFlag {
		if macFlag {
			res, err := resolve.ResolveMACs(matches)
			if err != nil {
				fatalf("%v", err)
			}
			if format == "json" {
				enc := json.NewEncoder(os.Stdout)
				enc.SetIndent("", "  ")
				must(enc.Encode(res))
				return
			}
			for _, row := range res {
				// one line per mapping
				for _, ip := range row.IPs {
					fmt.Printf("%s\t%s\n", row.MAC, ip)
				}
				if len(row.IPs) == 0 {
					fmt.Println(row.MAC)
				}
			}
			return
		}

		res := resolve.ReverseLookupAll(ctx, matches, timeout, jobs)
		if format == "json" {
			enc := json.NewEncoder(os.Stdout)
			enc.SetIndent("", "  ")
			must(enc.Encode(res))
			return
		}
		for _, row := range res {
			if row.Hostname != "" {
				fmt.Printf("%s\t%s\n", row.IP, row.Hostname)
			} else {
				fmt.Println(row.IP)
			}
		}
		return
	}

	// plain output
	if format == "json" {
		enc := json.NewEncoder(os.Stdout)
		enc.SetIndent("", "  ")
		must(enc.Encode(matches))
		return
	}
	if format == "csv" {
		// csv doesn't make much sense for plain extraction, but keep it consistent.
		w := csv.NewWriter(os.Stdout)
		must(w.Write([]string{"value"}))
		for _, m := range matches {
			must(w.Write([]string{m}))
		}
		w.Flush()
		if err := w.Error(); err != nil {
			fatalf("csv: %v", err)
		}
		return
	}
	for _, m := range matches {
		fmt.Println(m)
	}
}

func must(err error) {
	if err != nil {
		fatalf("%v", err)
	}
}

func fatalf(format string, args ...any) {
	fmt.Fprintf(os.Stderr, "ipgrep: "+format+"\n", args...)
	os.Exit(1)
}