topip/geo/topgeoip.go

package main

import (
	"bufio"
	"flag"
	"fmt"
	"log"
	"net"
	"os"
	"sort"
	"strings"

	"github.com/oschwald/geoip2-golang"
)

var (
	topk  = flag.Int("t", 10, "Number of top entries to show")
	logf  = flag.String("l", "/dev/stdin", "Path to the log files (reads from stdin by default)")
	cdbf  = flag.String("c", "/var/lib/GeoIP/GeoLite2-City.mmdb", "Path to the GeoIP City DB")
	adbf  = flag.String("a", "/var/lib/GeoIP/GeoLite2-ASN.mmdb", "Path to the GeoIP ASN DB")
	v4f   = flag.Bool("4", false, "IPv4 only")
	v6f   = flag.Bool("6", false, "IPv6 only")
	debug = flag.Bool("d", false, "Debug mode")
)

type kv struct {
	Key   string
	Value int
}

type geo struct {
	Country string
	ORG     string
	ASN     uint
}

func main() {
	flag.Parse()
	f, err := os.Open(*logf)
	if err != nil {
		fmt.Println(err)
		return
	}
	defer f.Close()

	if *debug {
		fmt.Println("Logs found...")
	}

	cityDB, err := geoip2.Open(*cdbf)
	if err != nil {
		log.Fatalf("opening City DB: %v", err)
	}
	defer cityDB.Close()

	if *debug {
		fmt.Println("City DB loaded...")
	}

	asnDB, err := geoip2.Open(*adbf)
	if err != nil {
		log.Fatalf("opening ASN DB: %v", err)
	}
	defer asnDB.Close()

	if *debug {
		fmt.Println("ASN DB loaded...")
		fmt.Println("Starting analysis...")
	}

	var addr []kv
	var countries []kv
	var asns []kv

	scanner := bufio.NewScanner(f)
	for scanner.Scan() {
		line := scanner.Text()
		ips := getIPs(line)
		for _, ip := range ips {
			geo, err := getGeo(ip, cityDB, asnDB)
			if err != nil {
				fmt.Println(err)
				return
			}
			addr = increment(addr, ip)
			countries = increment(countries, geo.Country)
			asns = increment(asns, geo.ORG)
		}
	}

	if err := scanner.Err(); err != nil {
		fmt.Println(err)
		return
	}

	if *debug {
		fmt.Println("Logs reading done, sorting results...")
	}

	sort.Slice(addr, func(i, j int) bool {
		return addr[i].Value > addr[j].Value
	})
	sort.Slice(countries, func(i, j int) bool {
		return countries[i].Value > countries[j].Value
	})
	sort.Slice(asns, func(i, j int) bool {
		return asns[i].Value > asns[j].Value
	})

	fmt.Printf("\033[1mResults:\033[0m\n\n")
	printStats(addr, "Top IP Addresses:")
	printStats(countries, "Top Countries:")
	printStats(asns, "Top Organizations:")
}

func getIPs(input string) []string {
	var ips []string
	words := strings.Fields(input)

	for _, word := range words {
		if ip := net.ParseIP(word); ip != nil {
			if !*v4f && !*v6f {
				ips = append(ips, ip.String())
			} else if *v4f && !strings.Contains(ip.String(), ".") {
				ips = append(ips, ip.String())
			} else if *v6f && !strings.Contains(ip.String(), ":") {
				ips = append(ips, ip.String())
			}
		}
	}

	return ips
}

func getGeo(ip string, cityDB, asnDB *geoip2.Reader) (geo, error) {
	var geo geo

	if asnDB != nil {
		record, err := asnDB.ASN(net.ParseIP(ip))
		if err != nil {
			geo.ORG = "Unknown"
			geo.ASN = 0
		}
		geo.ORG = record.AutonomousSystemOrganization
		geo.ASN = record.AutonomousSystemNumber
	}

	if cityDB != nil {
		record, err := cityDB.City(net.ParseIP(ip))
		if err != nil {
			geo.Country = "Unknown"
		}
		geo.Country = record.Country.Names["en"]
	}
	return geo, nil
}

func increment(data []kv, key string) []kv {
	if len(data) == 0 {
		return append(data, kv{key, 1})
	}
	for i, v := range data {
		if v.Key == key {
			data[i].Value++
			return data
		}
	}
	return append(data, kv{key, 1})
}

func printStats(stats []kv, title string) {
	var i int
	fmt.Printf("\033[1m%s\033[0m\n\n", title)
	for _, stat := range stats {
		fmt.Printf("  %d\t%s\n", stat.Value, stat.Key)
		i++
		if i >= *topk {
			break
		}
	}
	fmt.Println()
}