Cloudfish/mmdb_writer/mmdb_writer.go

package main

import (
	"encoding/json"
	"fmt"
	"math"
	"net"
	"os"

	"github.com/maxmind/mmdbwriter"
	"github.com/maxmind/mmdbwriter/mmdbtype"
)

func main() {
	if err := run(); err != nil {
		fmt.Fprintf(os.Stderr, "error: %v\n", err)
		os.Exit(1)
	}
}

func run() error {
	// Read the JSON object from stdin.
	//
	// Example:
	//
	// {
	//   "1.1.1.1": {
	//     "is_proxy": true,
	//     "score": 42
	//   }
	// }
	var input map[string]any

	decoder := json.NewDecoder(os.Stdin)
	decoder.UseNumber()

	if err := decoder.Decode(&input); err != nil {
		return fmt.Errorf("invalid JSON: %w", err)
	}

	// Create an IPv6 MMDB.
	//
	// IPv6 trees can contain both IPv6 and IPv4 networks.
	tree, err := mmdbwriter.New(mmdbwriter.Options{
		DatabaseType: "IP Classification",

		Description: map[string]string{
			"en": "Cloudfish",
		},

		IPVersion: 6,

		// Valid values are 24, 28, and 32.
		RecordSize: 28,
		
		IncludeReservedNetworks: true,
		
		DisableIPv4Aliasing: true,
	})
	if err != nil {
		return fmt.Errorf("creating MMDB: %w", err)
	}

	// Insert every IP/CIDR.
	for address, rawRecord := range input {
		network, err := parseNetwork(address)
		if err != nil {
			return fmt.Errorf("invalid IP/CIDR %q: %w", address, err)
		}

		record, err := toMMDBValue(rawRecord)
		if err != nil {
			return fmt.Errorf("invalid record for %q: %w", address, err)
		}

		if err := tree.Insert(network, record); err != nil {
			return fmt.Errorf("inserting %q: %w", address, err)
		}
	}

	// Write the binary MMDB to stdout.
	if _, err := tree.WriteTo(os.Stdout); err != nil {
		return fmt.Errorf("writing MMDB: %w", err)
	}

	return nil
}

// parseNetwork accepts:
//
//   1.1.1.1
//   1.0.0.0/16
//   2001:db8::1
//   2001:db8::/32
//
// A plain IPv4 address becomes /32.
// A plain IPv6 address becomes /128.
func parseNetwork(s string) (*net.IPNet, error) {
	// First try CIDR notation.
	if _, network, err := net.ParseCIDR(s); err == nil {
		return network, nil
	}

	// Then try a plain IP address.
	ip := net.ParseIP(s)
	if ip == nil {
		return nil, fmt.Errorf("not a valid IP address or CIDR")
	}

	// IPv4.
	if ip4 := ip.To4(); ip4 != nil {
		return &net.IPNet{
			IP:   ip4,
			Mask: net.CIDRMask(32, 32),
		}, nil
	}

	// IPv6.
	return &net.IPNet{
		IP:   ip,
		Mask: net.CIDRMask(128, 128),
	}, nil
}

// toMMDBValue recursively converts normal JSON values into
// MMDB values.
func toMMDBValue(v any) (mmdbtype.DataType, error) {
	switch x := v.(type) {

	case nil:
		return nil, nil

	case bool:
		return mmdbtype.Bool(x), nil

	case string:
		return mmdbtype.String(x), nil

	case json.Number:
		return jsonNumberToMMDB(x)

	case []any:
		array := make(mmdbtype.Slice, len(x))

		for i, item := range x {
			value, err := toMMDBValue(item)
			if err != nil {
				return nil, err
			}

			array[i] = value
		}

		return array, nil

	case map[string]any:
		m := make(mmdbtype.Map, len(x))

		for key, value := range x {
			converted, err := toMMDBValue(value)
			if err != nil {
				return nil, fmt.Errorf("%s: %w", key, err)
			}

			m[mmdbtype.String(key)] = converted
		}

		return m, nil

	default:
		return nil, fmt.Errorf("unsupported JSON type %T", v)
	}
}

// jsonNumberToMMDB converts a JSON number into the most appropriate
// MMDB numeric type.
//
// Integer values:
//
//   0 - 65535              -> Uint16
//   65536 - 4294967295     -> Uint32
//   4294967296 - MaxUint64 -> Uint64
//   negative values        -> Int32
//
// Non-integer values are stored as Float64.
func jsonNumberToMMDB(n json.Number) (mmdbtype.DataType, error) {
	s := n.String()

	// First try it as an integer.
	if i, err := n.Int64(); err == nil {
		// Positive integers.
		if i >= 0 {
			u := uint64(i)

			switch {
			case u <= math.MaxUint16:
				return mmdbtype.Uint16(u), nil

			case u <= math.MaxUint32:
				return mmdbtype.Uint32(u), nil

			default:
				return mmdbtype.Uint64(u), nil
			}
		}

		// Negative integers must fit in signed 32-bit MMDB type.
		if i >= math.MinInt32 {
			return mmdbtype.Int32(i), nil
		}

		return nil, fmt.Errorf(
			"integer %q is below the supported signed 32-bit range",
			s,
		)
	}

	// Not an integer, so try float64.
	f, err := n.Float64()
	if err != nil {
		return nil, fmt.Errorf("invalid number %q", s)
	}

	if math.IsNaN(f) || math.IsInf(f, 0) {
		return nil, fmt.Errorf("invalid floating-point number %q", s)
	}

	return mmdbtype.Float64(f), nil
}