~shabbyrobe/asnine

e4a200d24f48239d59baef08ee74f9fe707dcfa7 — Blake Williams 7 months ago 79b3eab
Update embedded data
3 files changed, 101 insertions(+), 3 deletions(-)

M asn/data/data-raw-table.gz
M asn/data/data-used-autnums.gz
M cmd/asnine/main.go
M asn/data/data-raw-table.gz => asn/data/data-raw-table.gz +0 -0
M asn/data/data-used-autnums.gz => asn/data/data-used-autnums.gz +0 -0
M cmd/asnine/main.go => cmd/asnine/main.go +101 -3
@@ 3,6 3,7 @@ package main
import (
	"bufio"
	"bytes"
	"compress/gzip"
	_ "embed"
	"encoding/json"
	"flag"


@@ 10,8 11,10 @@ import (
	"io"
	"log"
	"net"
	"net/http"
	"net/textproto"
	"os"
	"path/filepath"
	"sort"
	"strings"
	"time"


@@ 31,7 34,8 @@ var splash string
var usage = splash + "\n" + strings.TrimSpace(`
Usage: asnine [global] lookup <ip>
       asnine [global] scan [--fail-fast] (<file> | -)
       asnine [global] serve <addr>
       asnine [global] serve [<addr>]
       asnine [global] update [--force] <dest>

Global Flags:
  --asns <file>


@@ 49,7 53,9 @@ Commands:
          emit one jsonl record of results for each input.
  serve   Very primitive linewise TCP server, with the exact same semantics as
          scan (without --fail-fast). To keep the connection alive, send a newline
          regularly and you'll get a newline straight back.
          regularly and you'll get a newline straight back. Default addr is
          127.0.0.1:41014
  update  Download latest data
`)

func run() error {


@@ 80,6 86,9 @@ func run() error {
	case "serve":
		return cmdServe(conf, args)

	case "update":
		return cmdUpdate(conf, args)

	default:
		return fmt.Errorf("%s\n\nunknown command %q", usage, cmd)
	}


@@ 174,7 183,9 @@ func cmdServe(conf config, args []string) error {
	var readTimeout = 60 * time.Second
	var writeTimeout = 5 * time.Second

	if len(args) != 1 {
	if len(args) == 0 {
		args = []string{"127.0.0.1:41014"}
	} else if len(args) != 1 {
		return fmt.Errorf(usage)
	}



@@ 246,6 257,93 @@ func cmdServe(conf config, args []string) error {
	}
}

func cmdUpdate(conf config, args []string) error {
	var force bool
	fs := flag.NewFlagSet("", 0)
	fs.BoolVar(&force, "force", false, "Overwrite existing files")
	if err := fs.Parse(args); err != nil {
		return err
	}
	args = fs.Args()

	if len(args) != 1 {
		return fmt.Errorf(usage)
	}

	dest := args[0]
	if _, err := os.Stat(dest); err != nil {
		return err
	}

	fetch := func(u string) ([]byte, error) {
		rs, err := http.Get(u)
		if err != nil {
			return nil, err
		}
		defer rs.Body.Close()

		raw, err := io.ReadAll(rs.Body)
		if err != nil {
			return nil, err
		}

		var buf bytes.Buffer
		gzw, err := gzip.NewWriterLevel(&buf, gzip.BestCompression)
		if err != nil {
			return nil, err
		}
		defer gzw.Close()

		if _, err := gzw.Write(raw); err != nil {
			return nil, err
		}
		if err := gzw.Close(); err != nil {
			return nil, err
		}
		return buf.Bytes(), nil
	}

	type download struct {
		URL  string
		Dest string
	}

	downloads := []download{
		{asn.AddressMappingsURL, "data-raw-table.gz"},
		{asn.ReferenceURL, "data-used-autnums.gz"},
	}

	files := make([]string, len(downloads))
	for idx, dl := range downloads {
		file := filepath.Join(dest, dl.Dest)
		if !force {
			if _, err := os.Stat(file); err == nil {
				return fmt.Errorf("file %q exists", file)
			}
		}
		files[idx] = file
	}

	contents := make([][]byte, len(downloads))
	for idx, dl := range downloads {
		fmt.Printf("downloading %q\n", dl.URL)
		bts, err := fetch(dl.URL)
		if err != nil {
			return err
		}
		contents[idx] = bts
	}

	for idx, file := range files {
		fmt.Printf("writing %q\n", file)
		if err := os.WriteFile(file, contents[idx], 0o600); err != nil {
			return err
		}
	}

	return nil
}

func sortMappings(mappings []asn.Mapping) {
	sort.Slice(mappings, func(i, j int) bool {
		// Sort nil networks last