openvpn-mgt/vendor/github.com/pyke369/golang-support/prefixdb/cmd/prefixdb.go

368 lines
13 KiB
Go

package main
import (
"bufio"
"encoding/json"
"fmt"
"net"
"net/http"
"os"
"regexp"
"strconv"
"strings"
"time"
"github.com/pyke369/golang-support/prefixdb"
)
type LOCATION struct {
ContinentCode string
ContinentName string
CountryCode string
CountryName string
RegionCode string
RegionName string
StateCode string
StateName string
CityName string
TimeZone string
InEurope bool
}
var (
csvMatcher = regexp.MustCompile(`(?:,|\n|^)("(?:(?:"")*[^"]*)*"|[^",\n]*|(?:\n|$))`)
jsonMatcher = regexp.MustCompile(`^(\S+)(?:\s(\{.+?\}))?$`)
pfdb = prefixdb.New()
)
func size(input int) string {
if input < 1024*1024 {
return fmt.Sprintf("%.1fkB", float64(input)/1024)
}
return fmt.Sprintf("%.1fMB", float64(input)/(1024*1024))
}
func mkjson() {
for index := 3; index < len(os.Args); index++ {
count := 0
if handle, err := os.Open(os.Args[index]); err == nil {
reader := bufio.NewReader(handle)
last := time.Now()
start := last
for {
if line, err := reader.ReadString('\n'); err != nil {
break
} else {
if fields := jsonMatcher.FindStringSubmatch(strings.TrimSpace(line)); fields != nil {
if _, prefix, err := net.ParseCIDR(fields[1]); err == nil {
data := map[string]interface{}{}
json.Unmarshal([]byte(fields[2]), &data)
pfdb.Add(*prefix, data, [][]string{[]string{"key1", "key2"}})
count++
}
}
}
if now := time.Now(); now.Sub(last) >= 250*time.Millisecond {
last = now
fmt.Fprintf(os.Stderr, "\r- adding prefixes [%s] %d", os.Args[index], count)
}
}
handle.Close()
fmt.Fprintf(os.Stderr, "\r- added prefixes [%s] (%.3fs - %d entries)\n", os.Args[index], float64(time.Now().Sub(start))/float64(time.Second), count)
}
}
start := time.Now()
description := ""
if index := strings.Index(os.Args[2], "@"); index > 0 {
description = os.Args[2][index+1:]
os.Args[2] = os.Args[2][:index]
}
fmt.Fprintf(os.Stderr, "\r- saving database [%s]... ", os.Args[2])
if _, err := pfdb.Save(os.Args[2], description); err != nil {
fmt.Fprintf(os.Stderr, "\r- saving database [%s] failed (%v)\n", os.Args[2], err)
} else {
fmt.Fprintf(os.Stderr, "\r- saved database [%s] (%.3fs - total[%s] strings[%s] numbers[%s] pairs[%s] clusters[%s] maps[%s] nodes[%s])\n",
os.Args[2], float64(time.Now().Sub(start))/float64(time.Second), size(pfdb.Total), size(pfdb.Strings[0]),
size(pfdb.Numbers[0]), size(pfdb.Pairs[0]), size(pfdb.Clusters[0]), size(pfdb.Maps[0]), size(pfdb.Nodes[0]),
)
}
}
func mkcity() {
locations := map[int]*LOCATION{}
if handle, err := os.Open(os.Args[3]); err == nil {
reader := bufio.NewReader(handle)
last := time.Now()
start := last
for {
if line, err := reader.ReadString('\n'); err != nil {
break
} else {
if fields := csvMatcher.FindAllStringSubmatch(strings.TrimSpace(line), -1); len(fields) == 14 {
for index := 0; index < len(fields); index++ {
fields[index][1] = strings.Trim(fields[index][1], `"`)
}
if id, err := strconv.Atoi(fields[0][1]); err == nil {
locations[id] = &LOCATION{
ContinentCode: fields[2][1],
ContinentName: fields[3][1],
CountryCode: fields[4][1],
CountryName: fields[5][1],
RegionCode: fields[6][1],
RegionName: fields[7][1],
StateCode: fields[8][1],
StateName: fields[9][1],
CityName: fields[10][1],
TimeZone: fields[12][1],
InEurope: fields[13][1] == "1",
}
}
}
}
if now := time.Now(); now.Sub(last) >= 250*time.Millisecond {
last = now
fmt.Fprintf(os.Stderr, "\r- loading locations [%s] %d", os.Args[3], len(locations))
}
}
handle.Close()
fmt.Fprintf(os.Stderr, "\r- loaded locations [%s] (%.3fs - %d entries)\n", os.Args[3], float64(time.Now().Sub(start))/float64(time.Second), len(locations))
}
clusters := [][]string{
[]string{"continent_code", "continent_name", "country_code", "country_name", "region_code", "region_name", "state_code", "state_name", "timezone", "in_europe"},
[]string{"city_name", "postal_code", "latitude", "longitude"},
}
for index := 4; index < len(os.Args); index++ {
count := 0
if handle, err := os.Open(os.Args[index]); err == nil {
reader := bufio.NewReader(handle)
last := time.Now()
start := last
for {
if line, err := reader.ReadString('\n'); err != nil {
break
} else {
if fields := strings.Split(strings.TrimSpace(line), ","); len(fields) == 10 {
id := 0
if id, _ = strconv.Atoi(fields[1]); id == 0 {
id, _ = strconv.Atoi(fields[2])
}
if id != 0 && locations[id] != nil {
if _, prefix, err := net.ParseCIDR(fields[0]); err == nil {
latitude, _ := strconv.ParseFloat(fields[7], 64)
longitude, _ := strconv.ParseFloat(fields[8], 64)
pfdb.Add(*prefix, map[string]interface{}{
"continent_code": locations[id].ContinentCode,
"continent_name": locations[id].ContinentName,
"country_code": locations[id].CountryCode,
"country_name": locations[id].CountryName,
"region_code": locations[id].RegionCode,
"region_name": locations[id].RegionName,
"state_code": locations[id].StateCode,
"state_name": locations[id].StateName,
"city_name": locations[id].CityName,
"in_europe": locations[id].InEurope,
"timezone": locations[id].TimeZone,
"postal_code": fields[6],
"latitude": latitude,
"longitude": longitude,
}, clusters)
count++
}
}
}
}
if now := time.Now(); now.Sub(last) >= 250*time.Millisecond {
last = now
fmt.Fprintf(os.Stderr, "\r- adding prefixes [%s] %d", os.Args[index], count)
}
}
handle.Close()
fmt.Fprintf(os.Stderr, "\r- added prefixes [%s] (%.3fs - %d entries)\n", os.Args[index], float64(time.Now().Sub(start))/float64(time.Second), count)
}
}
start := time.Now()
description := ""
if index := strings.Index(os.Args[2], "@"); index > 0 {
description = os.Args[2][index+1:]
os.Args[2] = os.Args[2][:index]
}
fmt.Fprintf(os.Stderr, "\r- saving database [%s]... ", os.Args[2])
if _, err := pfdb.Save(os.Args[2], description); err != nil {
fmt.Fprintf(os.Stderr, "\r- saving database [%s] failed (%v)\n", os.Args[2], err)
} else {
fmt.Fprintf(os.Stderr, "\r- saved database [%s] (%.3fs - total[%s] strings[%s] numbers[%s] pairs[%s] clusters[%s] maps[%s] nodes[%s])\n",
os.Args[2], float64(time.Now().Sub(start))/float64(time.Second), size(pfdb.Total), size(pfdb.Strings[0]),
size(pfdb.Numbers[0]), size(pfdb.Pairs[0]), size(pfdb.Clusters[0]), size(pfdb.Maps[0]), size(pfdb.Nodes[0]),
)
}
}
func mkasn() {
for index := 3; index < len(os.Args); index++ {
count := 0
if handle, err := os.Open(os.Args[index]); err == nil {
reader := bufio.NewReader(handle)
last := time.Now()
start := last
for {
if line, err := reader.ReadString('\n'); err != nil {
break
} else {
if fields := csvMatcher.FindAllStringSubmatch(strings.TrimSpace(line), -1); len(fields) == 3 {
for index := 0; index < len(fields); index++ {
fields[index][1] = strings.Trim(fields[index][1], `"`)
}
if asnum, _ := strconv.Atoi(fields[1][1]); asnum != 0 {
if _, prefix, err := net.ParseCIDR(fields[0][1]); err == nil {
pfdb.Add(*prefix, map[string]interface{}{
"as_number": fmt.Sprintf("AS%d", asnum),
"as_name": fields[2][1],
}, nil)
count++
}
}
}
}
if now := time.Now(); now.Sub(last) >= 250*time.Millisecond {
last = now
fmt.Fprintf(os.Stderr, "\r- adding prefixes [%s] %d", os.Args[index], count)
}
}
handle.Close()
fmt.Fprintf(os.Stderr, "\r- added prefixes [%s] (%.3fs - %d entries)\n", os.Args[index], float64(time.Now().Sub(start))/float64(time.Second), count)
}
}
start := time.Now()
description := ""
if index := strings.Index(os.Args[2], "@"); index > 0 {
description = os.Args[2][index+1:]
os.Args[2] = os.Args[2][:index]
}
fmt.Fprintf(os.Stderr, "\r- saving database [%s]... ", os.Args[2])
if _, err := pfdb.Save(os.Args[2], description); err != nil {
fmt.Fprintf(os.Stderr, "\r- saving database [%s] failed (%v)\n", os.Args[2], err)
} else {
fmt.Fprintf(os.Stderr, "\r- saved database [%s] (%.3fs - total[%s] strings[%s] numbers[%s] pairs[%s] clusters[%s] maps[%s] nodes[%s])\n",
os.Args[2], float64(time.Now().Sub(start))/float64(time.Second), size(pfdb.Total), size(pfdb.Strings[0]),
size(pfdb.Numbers[0]), size(pfdb.Pairs[0]), size(pfdb.Clusters[0]), size(pfdb.Maps[0]), size(pfdb.Nodes[0]),
)
}
}
func lookup() {
databases := []*prefixdb.PrefixDB{}
for index := 2; index < len(os.Args); index++ {
if strings.HasSuffix(os.Args[index], `.pfdb`) {
fmt.Fprintf(os.Stderr, "\r- loading database [%s]...", os.Args[index])
database := prefixdb.New()
if err := database.Load(os.Args[index]); err == nil {
fmt.Fprintf(os.Stderr, "\r- loaded database [%s] (total[%s] version[%d.%d.%d] description[%s])\n",
os.Args[index], size(database.Total), (database.Version>>16)&0xff, (database.Version>>8)&0xff, database.Version&0xff, database.Description)
databases = append(databases, database)
} else {
fmt.Fprintf(os.Stderr, "\r- loading database [%s] failed (%v)\n", os.Args[index], err)
}
} else {
if ip := net.ParseIP(os.Args[index]); ip == nil {
fmt.Fprintf(os.Stderr, "- lookup [%s] failed (not a valid IP address)", os.Args[index])
} else {
fmt.Fprintf(os.Stderr, "- lookup [%s] ", os.Args[index])
lookup := map[string]interface{}{}
for _, database := range databases {
lookup, _ = database.Lookup(ip, lookup)
}
data, _ := json.Marshal(lookup)
fmt.Printf("%s\n", data)
}
}
}
}
func server() {
databases := []*prefixdb.PrefixDB{}
for index := 3; index < len(os.Args); index++ {
fmt.Fprintf(os.Stderr, "\r- loading database [%s]...", os.Args[index])
database := prefixdb.New()
if err := database.Load(os.Args[index]); err == nil {
fmt.Fprintf(os.Stderr, "\r- loaded database [%s] (total[%s] version[%d.%d.%d] description[%s])\n",
os.Args[index], size(database.Total), (database.Version>>16)&0xff, (database.Version>>8)&0xff, database.Version&0xff, database.Description)
databases = append(databases, database)
} else {
fmt.Fprintf(os.Stderr, "\r- loading database [%s] failed (%v)\n", os.Args[index], err)
}
}
http.HandleFunc("/", func(response http.ResponseWriter, request *http.Request) {
response.Header().Set("Content-Type", "application/json")
remote, _, _ := net.SplitHostPort(request.RemoteAddr)
if value := request.Header.Get("X-Forwarded-For"); value != "" {
remote = strings.Split(value, ",")[0]
}
parameters := request.URL.Query()
if value := parameters.Get("remote"); value != "" {
remote = value
}
lookup := map[string]interface{}{}
if ip := net.ParseIP(remote); ip != nil {
lookup["ip"] = fmt.Sprintf("%s", ip)
for _, database := range databases {
lookup, _ = database.Lookup(ip, lookup)
}
data, _ := json.Marshal(lookup)
response.Write(data)
}
})
parts := strings.Split(os.Args[2], ",")
server := &http.Server{Addr: strings.TrimLeft(parts[0], "*"), ReadTimeout: 10 * time.Second, WriteTimeout: 10 * time.Second}
fmt.Fprintf(os.Stderr, "\r- listening to [%s]\n", os.Args[2])
if len(parts) > 1 {
server.ListenAndServeTLS(parts[1], parts[2])
} else {
server.ListenAndServe()
}
}
func usage(status int) {
fmt.Fprintf(os.Stderr, "usage: prefixdb <action> [parameters...]\n\n"+
"help show this help screen\n"+
"json <database[@<description>]> <JSON prefixes>... build database from generic JSON-formatted prefixes lists\n"+
"city <database[@<description>]> <CSV locations> <CSV prefixes>... build database from MaxMind GeoIP2 cities lists\n"+
"asn <database[@<description>]> <CSV prefixes>... build database from MaxMind GeoLite2 asnums lists\n"+
"lookup <database>... <address>... lookup entries in databases\n"+
"server <bind address> <database>... spawn an HTTP(S) server for entries lookup\n")
os.Exit(status)
}
func main() {
if len(os.Args) < 2 {
usage(1)
}
switch os.Args[1] {
case "help":
usage(0)
case "json":
if len(os.Args) < 3 {
usage(1)
}
mkjson()
case "city":
if len(os.Args) < 5 {
usage(1)
}
mkcity()
case "asn":
if len(os.Args) < 3 {
usage(1)
}
mkasn()
case "lookup":
lookup()
case "server":
server()
default:
usage(2)
}
}