lgregmapper/regclient.go

//////////////////////////////////////////////////////////////////////////
// lgregmapper, a small utility to provide DN42 registry data to bird-lg
//////////////////////////////////////////////////////////////////////////

package main

//////////////////////////////////////////////////////////////////////////

import (
	"fmt"
	log "github.com/sirupsen/logrus"
	"net/http"
	"strings"
	"sync"
	//	"net"
	"encoding/json"
	//	"io/ioutil"
	"time"
)

//////////////////////////////////////////////////////////////////////////

type RegClient struct {
	done     chan bool
	wait     *sync.WaitGroup
	endpoint string
	client   *http.Client
	ASN      map[string][]byte
}

type APIAttributeResponse map[string]map[string][]string

type ASNData struct {
	mntner   []string
	descr    string
	prefixes []string
}

//////////////////////////////////////////////////////////////////////////
// create a new RegClient object

func NewRegClient(apiAddr string) *RegClient {

	rc := &RegClient{
		done:     make(chan bool),
		wait:     &sync.WaitGroup{},
		client:   &http.Client{Timeout: 20 * time.Second},
		endpoint: apiAddr,
		ASN:      nil,
	}

	rc.wait.Add(1)
	return rc
}

//////////////////////////////////////////////////////////////////////////
// regularly update the registry

func (rc *RegClient) Start(interval time.Duration) {
	defer rc.wait.Done()

	ticker := time.NewTicker(interval)
	defer ticker.Stop()

	for {
		select {
		case <-rc.done:
			return

		case <-ticker.C:
			rc.Update()
		}
	}
}

//////////////////////////////////////////////////////////////////////////
// shutdown the client

func (rc *RegClient) Shutdown() {
	close(rc.done)
	rc.wait.Wait()
	log.Info("Registry client shutdown")
}

//////////////////////////////////////////////////////////////////////////
// fetch data from the registry API

func (rc *RegClient) queryAPI(query string, target interface{}) error {

	query = rc.endpoint + "/api/registry/" + query

	log.WithFields(log.Fields{
		"Query": query,
	}).Debug("Querying API")

	response, err := rc.client.Get(query)
	if err != nil {
		log.WithFields(log.Fields{
			"Error": err,
			"Query": query,
		}).Error("Unable to query API")
		return err
	}
	defer response.Body.Close()

	if err := json.NewDecoder(response.Body).Decode(target); err != nil {
		log.WithFields(log.Fields{
			"Error": err,
		}).Error("Failed to decode JSON")
		return err
	}
	return nil
}

//////////////////////////////////////////////////////////////////////////
// perform an update

func (rc *RegClient) Update() {

	// query the API to collect ASN maintainers
	mnts := make(APIAttributeResponse)
	if err := rc.queryAPI("aut-num/*/mnt-by?raw", &mnts); err != nil {
		return
	}

	// and ASN descriptions
	descr := make(APIAttributeResponse)
	if err := rc.queryAPI("aut-num/*/descr?raw", &descr); err != nil {
		return
	}

	// and query again to collect origin data
	origins := make(APIAttributeResponse)
	if err := rc.queryAPI("*route/*/origin?raw", &origins); err != nil {
		return
	}

	// normalise the returned data in to a struct
	asnData := make(map[string]*ASNData)
	for asnpath, adata := range mnts {

		asn := strings.TrimPrefix(asnpath, "aut-num/")
		var description string

		dmap := descr[asnpath]
		if dmap != nil && dmap["descr"] != nil && len(dmap["descr"]) > 0 {
			description = dmap["descr"][0]
		}

		asnData[asn] = &ASNData{
			mntner:   adata["mnt-by"],
			descr:    description,
			prefixes: nil,
		}
	}

	// and add the origin data
	for prefix, odata := range origins {
		ix := strings.IndexRune(prefix, '/')
		prefix = strings.Replace(prefix[ix+1:], "_", "/", 1)
		for _, origin := range odata["origin"] {
			if asnd := asnData[origin]; asnd != nil {
				asnd.prefixes = append(asnd.prefixes, prefix)
			}
		}
	}

	// finally, turn the collated data in to a string suitable
	// for returning from the fake memcache server
	asns := make(map[string][]byte)
	for asn, adata := range asnData {

		asn = "lg_" + strings.TrimPrefix(asn, "AS")

		plist := strings.Join(adata.prefixes, "\r")
		if plist != "" {
			// add spacer if not empty
			plist = "\r" + plist
		}

		if adata.descr != "" {
			adata.descr += "\r"
		}

		data := fmt.Sprintf("%s\r%s%s",
			strings.Join(adata.mntner, "\r"),
			adata.descr, plist)

		asns[asn] = []byte(fmt.Sprintf("VALUE %s 0 %d\r\n%s\r\nEND\r\n",
			asn, len(data), data))
	}

	log.WithFields(log.Fields{
		"ASN Count": len(asns),
	}).Debug("Updated ASN data")

	rc.ASN = asns
}

//////////////////////////////////////////////////////////////////////////
// end of code