port ipgrep to golang
This commit is contained in:
Tobias Kessels
BIN
tools/go/goipgrep/ipgrep
Executable file
BIN
tools/go/goipgrep/ipgrep
Executable file
Binary file not shown.
376
tools/go/goipgrep/ipgrep.go
Normal file
376
tools/go/goipgrep/ipgrep.go
Normal file
@@ -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)
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user