maxmind_isp_origin.go

package main

import (
	"bytes"
	"errors"
	"fmt"
	"net"
	"strings"
	"text/template"

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

func MaxMindHeader(valueDescription string, dbMetadata maxminddb.Metadata) (string, error) {
	// TODO can this be a global const or something...
	tmpl, err := template.New("maxmind").Parse(`#
# MaxMind IP to {{.ValueType}} mapping
# 
# Database Info
#   Build Epoch:  {{.BuildEpoch}}
#   Description:  {{.Description}}
#   DatabaseType: {{.DatabaseType}}
#   NodeCount:    {{.NodeCount}}
#
# Generated by haproxy-mapper at {{.Now}}
#
`,
	)
	if err != nil {
		return "", err
	}

	description := "Not Available"
	description, _ = dbMetadata.Description["en"]

	context := struct {
		BuildEpoch   string
		DatabaseType string
		Description  string
		NodeCount    uint
		Now          string
		ValueType    string
	}{
		BuildEpoch:   EpochToIso8601(dbMetadata.BuildEpoch),
		DatabaseType: dbMetadata.DatabaseType,
		Description:  description,
		NodeCount:    dbMetadata.NodeCount,
		Now:          NowIso8610(),
		ValueType:    valueDescription,
	}

	var rendered bytes.Buffer
	if err = tmpl.Execute(&rendered, context); err != nil {
		return "", err
	}

	return rendered.String(), nil
}

type recordIsp struct {
	Asn int    `maxminddb:"autonomous_system_number"`
	Isp string `maxminddb:"isp"`
}

type MaxMindIspOrigin struct {
	Filename    string
	HaveIspData bool
	db          *maxminddb.Reader
	asn         Emitter
	isp         Emitter
}

func MaxMindIspOriginCreate(filename string) (*MaxMindIspOrigin, error) {
	// TODO: look at db.Metadata.Languages to make sure we have en and/or pick a language
	db, err := maxminddb.Open(filename)
	if err != nil {
		return nil, err
	}
	haveIspData := false
	if strings.HasSuffix(db.Metadata.DatabaseType, "ISP") {
		haveIspData = true
	} else if !strings.HasSuffix(db.Metadata.DatabaseType, "ASN") {
		return nil, errors.New("provided database is not MaxMind ISP or ASN")
	}

	return &MaxMindIspOrigin{
		Filename:    filename,
		HaveIspData: haveIspData,
		db:          db,
		asn: Emitter{
			id: "maxmind.asn",
		},
		isp: Emitter{
			id: "maxmind.isp",
		},
	}, nil
}

func (m *MaxMindIspOrigin) AddAsnReceiver(receiver Receiver) {
	m.asn.AddReceiver(receiver)
}

func (m *MaxMindIspOrigin) AddIspReceiver(receiver Receiver) {
	m.isp.AddReceiver(receiver)
}

func (m *MaxMindIspOrigin) headers() error {
	asn, err := MaxMindHeader("Autonomous System Number (ASN)", m.db.Metadata)
	if err != nil {
		return err
	}

	header := Header{
		general: asn,
		columns: "# cidr asn\n",
	}
	if err = m.asn.Header(header); err != nil {
		return err
	}

	isp, err := MaxMindHeader("Internet Service Provider (ISP)", m.db.Metadata)
	if err != nil {
		return err
	}

	header = Header{
		general: isp,
		columns: "# cidr isp\n",
	}
	return m.isp.Header(header)
}

func (m *MaxMindIspOrigin) Run(ipv4Only bool) error {
	if err := m.headers(); err != nil {
		return err
	}

	cidr := "::/0"
	if ipv4Only {
		cidr = "0.0.0.0/0"
	}
	_, network, err := net.ParseCIDR(cidr)
	if err != nil {
		return err
	}
	networks := m.db.NetworksWithin(network)
	if !ipv4Only {
		// We only want Ipv4 addresses once, if we call this when iterating
		// over 0.0.0.0/8 we get nothing for some reason
		maxminddb.SkipAliasedNetworks(networks)
	}

	for networks.Next() {
		record := recordIsp{}
		net, err := networks.Network(&record)
		if err != nil {
			return err
		}

		if record.Asn != 0 {
			asn := fmt.Sprintf("AS%d", record.Asn)

			err = m.asn.Emit(BlockCreate(net, &asn))
			if err != nil {
				return err
			}
		}

		if record.Isp != "" {
			err = m.isp.Emit(BlockCreate(net, &record.Isp))
			if err != nil {
				return err
			}
		}
	}

	m.asn.Done()
	m.isp.Done()

	return nil
}