port ipgrep to golang

tools/go/goipgrep/ipgrep.go

@@ -0,0 +1,376 @@
+// ipgrep.go
+package main
+
+import (
+	"bytes"
+	"encoding/csv"
+	"encoding/json"
+	"flag"
+	"fmt"
+	"io"
+	"io/ioutil"
+	"net/http"
+	"os"
+	"os/exec"
+	"regexp"
+	"sort"
+	"sync"
+	"time"
+)
+
+// IPInfo holds the data we want from ipinfo.io
+type IPInfo struct {
+	IP       string `json:"ip"`
+	Hostname string `json:"hostname"`
+	City     string `json:"city"`
+	Region   string `json:"region"`
+	Country  string `json:"country"`
+	Org      string `json:"org"`
+}
+
+func main() {
+	// Command-line flags.
+	var (
+		sortFlag    bool
+		uniqFlag    bool
+		macFlag     bool
+		pingable    bool
+		resolveFlag bool
+		lookupFlag  bool
+		fileName    string
+	)
+	flag.BoolVar(&uniqFlag, "u", false, "only show uniq IPs/MACs (implies -s)")
+	flag.BoolVar(&sortFlag, "s", false, "sort output")
+	flag.BoolVar(&macFlag, "m", false, "grep MAC-IDs instead of IPs")
+	flag.BoolVar(&pingable, "p", false, "only show 'pingable' entries (MACs still beta)")
+	flag.BoolVar(&resolveFlag, "r", false, "resolve (uses host for ip and arping for mac)")
+	flag.BoolVar(&lookupFlag, "l", false, "lookup ip info using ipinfo.io and output CSV: ip,country,region,city,org,hostname")
+	flag.StringVar(&fileName, "f", "", "input file")
+	flag.Usage = func() {
+		fmt.Fprintf(os.Stderr, "Usage: %s [-u] [-s] [-m] [-p] [-r] [-l] [-f filename] [file...]\n", os.Args[0])
+		flag.PrintDefaults()
+	}
+	flag.Parse()
+
+	// If pingable is set, force sorting and uniqueness.
+	if pingable || lookupFlag {
+		sortFlag = true
+		uniqFlag = true
+	}
+
+	if lookupFlag && macFlag {
+		fmt.Fprintln(os.Stderr, "Lookup mode (-l) only works for IP addresses, not MAC addresses.")
+		os.Exit(1)
+	}
+
+	// Regular expressions for IPs or MACs.
+	var pattern string
+	if macFlag {
+		// Supports MAC formats: xx:xx:xx:xx:xx:xx or xxxx.xxxx.xxxx
+		pattern = `(([a-fA-F0-9]{2}[:-]){5}[a-fA-F0-9]{2})|([a-fA-F0-9]{4}\.[a-fA-F0-9]{4}\.[a-fA-F0-9]{4})`
+	} else {
+		// Matches valid IPv4 addresses.
+		pattern = `(((25[0-5])|(2[0-4][0-9])|([0-1]?\d?\d))\.){3}((25[0-5])|(2[0-4][0-9])|([0-1]?\d?\d))`
+	}
+	re, err := regexp.Compile(pattern)
+	if err != nil {
+		fmt.Fprintf(os.Stderr, "Error compiling regex: %v\n", err)
+		os.Exit(1)
+	}
+
+	// Read input from -f file, extra args, or stdin.
+	var inputData []byte
+	if fileName != "" {
+		inputData, err = ioutil.ReadFile(fileName)
+		if err != nil {
+			fmt.Fprintf(os.Stderr, "Error reading file %s: %v\n", fileName, err)
+			os.Exit(1)
+		}
+	} else if flag.NArg() > 0 {
+		var buf bytes.Buffer
+		for _, fname := range flag.Args() {
+			data, err := ioutil.ReadFile(fname)
+			if err != nil {
+				fmt.Fprintf(os.Stderr, "Error reading file %s: %v\n", fname, err)
+				continue
+			}
+			buf.Write(data)
+			buf.WriteByte('\n')
+		}
+		inputData = buf.Bytes()
+	} else {
+		inputData, err = io.ReadAll(os.Stdin)
+		if err != nil {
+			fmt.Fprintf(os.Stderr, "Error reading stdin: %v\n", err)
+			os.Exit(1)
+		}
+	}
+
+	// Filter matches using the regex.
+	matches := re.FindAllString(string(inputData), -1)
+	if matches == nil {
+		os.Exit(0)
+	}
+
+	if sortFlag {
+		sort.Strings(matches)
+	}
+	if uniqFlag {
+		matches = unique(matches)
+	}
+
+	if pingable {
+		matches = filterPingable(matches, macFlag)
+		if sortFlag {
+			sort.Strings(matches)
+		}
+	}
+
+	// If lookup flag is set, perform ipinfo.io lookups with caching.
+	if lookupFlag {
+		cache := loadCache()
+		var cacheMu sync.Mutex
+		results := lookupIPInfo(matches, cache, &cacheMu)
+
+		// Sort the results by IP.
+		sort.Slice(results, func(i, j int) bool {
+			return results[i].IP < results[j].IP
+		})
+
+		// Save the updated cache.
+		saveCache(cache)
+
+		// Output CSV using csv.Writer for proper CSV formatting.
+		w := csv.NewWriter(os.Stdout)
+		// Write header.
+		if err := w.Write([]string{"ip", "country", "region", "city", "org", "hostname"}); err != nil {
+			fmt.Fprintf(os.Stderr, "Error writing CSV header: %v\n", err)
+			os.Exit(1)
+		}
+		for _, info := range results {
+			record := []string{
+				info.IP,
+				info.Country,
+				info.Region,
+				info.City,
+				info.Org,
+				info.Hostname,
+			}
+			if err := w.Write(record); err != nil {
+				fmt.Fprintf(os.Stderr, "Error writing CSV record: %v\n", err)
+				os.Exit(1)
+			}
+		}
+		w.Flush()
+		if err := w.Error(); err != nil {
+			fmt.Fprintf(os.Stderr, "Error flushing CSV data: %v\n", err)
+		}
+		return
+	}
+
+	// If resolve flag is set, perform resolution.
+	if resolveFlag {
+		results := resolveEntries(matches, macFlag)
+		for _, r := range results {
+			fmt.Print(r)
+			if !macFlag {
+				fmt.Println()
+			}
+		}
+		return
+	}
+
+	// Otherwise, just output the matches.
+	for _, m := range matches {
+		fmt.Println(m)
+	}
+}
+
+// unique removes duplicate strings from a slice.
+func unique(input []string) []string {
+	seen := make(map[string]struct{})
+	var result []string
+	for _, s := range input {
+		if _, ok := seen[s]; !ok {
+			seen[s] = struct{}{}
+			result = append(result, s)
+		}
+	}
+	return result
+}
+
+// filterPingable runs ping (or arping) concurrently on each entry.
+func filterPingable(entries []string, mac bool) []string {
+	var wg sync.WaitGroup
+	var mu sync.Mutex
+	var result []string
+
+	for _, entry := range entries {
+		wg.Add(1)
+		go func(e string) {
+			defer wg.Done()
+			if isPingable(e, mac) {
+				mu.Lock()
+				result = append(result, e)
+				mu.Unlock()
+			}
+		}(entry)
+	}
+	wg.Wait()
+	return result
+}
+
+// isPingable tests if an entry is reachable using ping (or arping for MACs).
+func isPingable(entry string, mac bool) bool {
+	var cmd *exec.Cmd
+	if mac {
+		cmd = exec.Command("arping", "-c", "1", "-w", "5000000", entry)
+	} else {
+		cmd = exec.Command("ping", "-c", "1", "-w", "1", entry)
+	}
+	cmd.Stdout = nil
+	cmd.Stderr = nil
+	err := cmd.Run()
+	return err == nil
+}
+
+// resolveEntries performs resolution via external commands.
+func resolveEntries(entries []string, mac bool) []string {
+	var wg sync.WaitGroup
+	results := make([]string, len(entries))
+
+	for i, entry := range entries {
+		wg.Add(1)
+		go func(i int, e string) {
+			defer wg.Done()
+			if mac {
+				cmd := exec.Command("arping", "-q", "-c", "1", "-w", "5000000", e)
+				if err := cmd.Run(); err == nil {
+					cmd2 := exec.Command("arping", "-c", "1", e)
+					out, err := cmd2.CombinedOutput()
+					if err == nil {
+						results[i] = string(out)
+					} else {
+						results[i] = e + "\n"
+					}
+				} else {
+					results[i] = e + "\n"
+				}
+			} else {
+				cmd := exec.Command("host", e)
+				out, err := cmd.CombinedOutput()
+				if err == nil {
+					// Extract the hostname via regex (similar to grep -Po '(?<=pointer ).*')
+					reHost := regexp.MustCompile(`(?i)(?<=pointer\s)(\S+)`)
+					match := reHost.FindString(string(out))
+					if match != "" {
+						results[i] = fmt.Sprintf("%s  %s", e, match)
+					} else {
+						results[i] = e
+					}
+				} else {
+					results[i] = e
+				}
+			}
+		}(i, entry)
+	}
+	wg.Wait()
+	return results
+}
+
+// lookupIPInfo queries ipinfo.io for each IP concurrently,
+// checking a local cache before going to the network.
+// It returns a slice of IPInfo.
+func lookupIPInfo(entries []string, cache map[string]IPInfo, cacheMu *sync.Mutex) []IPInfo {
+	var wg sync.WaitGroup
+	var mu sync.Mutex
+	var results []IPInfo
+
+	client := &http.Client{
+		Timeout: 5 * time.Second,
+	}
+
+	for _, ip := range entries {
+		wg.Add(1)
+		go func(ip string) {
+			defer wg.Done()
+			// Check cache first.
+			cacheMu.Lock()
+			info, found := cache[ip]
+			cacheMu.Unlock()
+			if found {
+				mu.Lock()
+				results = append(results, info)
+				mu.Unlock()
+				return
+			}
+
+			// Not in cache; perform HTTP lookup.
+			url := fmt.Sprintf("https://ipinfo.io/%s", ip)
+			resp, err := client.Get(url)
+			if err != nil {
+				return
+			}
+			defer resp.Body.Close()
+			body, err := io.ReadAll(resp.Body)
+			if err != nil {
+				return
+			}
+			var newInfo IPInfo
+			if err := json.Unmarshal(body, &newInfo); err != nil {
+				return
+			}
+			// Only add valid responses.
+			if newInfo.IP == "" {
+				return
+			}
+			// Update cache.
+			cacheMu.Lock()
+			cache[ip] = newInfo
+			cacheMu.Unlock()
+			mu.Lock()
+			results = append(results, newInfo)
+			mu.Unlock()
+		}(ip)
+	}
+	wg.Wait()
+	return results
+}
+
+// loadCache reads the cache from ~/.ipgrep.db (if present)
+// and returns it as a map[string]IPInfo.
+func loadCache() map[string]IPInfo {
+	home, err := os.UserHomeDir()
+	if err != nil {
+		return make(map[string]IPInfo)
+	}
+	cachePath := home + "/.ipgrep.db"
+	data, err := ioutil.ReadFile(cachePath)
+	if err != nil {
+		// File doesn't exist or can't be read, start with an empty cache.
+		return make(map[string]IPInfo)
+	}
+	var cache map[string]IPInfo
+	if err := json.Unmarshal(data, &cache); err != nil {
+		// If unmarshal fails, use an empty cache.
+		return make(map[string]IPInfo)
+	}
+	return cache
+}
+
+// saveCache writes the cache map to ~/.ipgrep.db.
+func saveCache(cache map[string]IPInfo) {
+	home, err := os.UserHomeDir()
+	if err != nil {
+		return
+	}
+	cachePath := home + "/.ipgrep.db"
+	data, err := json.MarshalIndent(cache, "", "  ")
+	if err != nil {
+		fmt.Fprintf(os.Stderr, "Error marshaling cache: %v\n", err)
+		return
+	}
+	if err := ioutil.WriteFile(cachePath, data, 0644); err != nil {
+		fmt.Fprintf(os.Stderr, "Error writing cache file: %v\n", err)
+	}
+}