From 7794ec7bcb6a3a149bb8121ceb5ac2577a05006e Mon Sep 17 00:00:00 2001
From: Lan Tian <xuyh0120@outlook.com>
Date: Fri, 27 Mar 2020 12:52:09 +0800
Subject: [PATCH] Add query AS info via DNS interface

---
 frontend/bgpmap.go | 18 ++++++++++++++----
 frontend/main.go   | 43 +++++++++++++++++++++++++++----------------
 2 files changed, 41 insertions(+), 20 deletions(-)

diff --git a/frontend/bgpmap.go b/frontend/bgpmap.go
index fc8e873..0ace7e9 100644
--- a/frontend/bgpmap.go
+++ b/frontend/bgpmap.go
@@ -1,9 +1,19 @@
 package main
 
 import (
+	"fmt"
+	"net"
 	"strings"
 )
 
+func getASNRepresentation(asn string) string {
+	if records, err := net.LookupTXT(fmt.Sprintf("AS%s.%s", asn, setting.dnsInterface)); err != nil {
+		return fmt.Sprintf("AS%s", asn)
+	} else {
+		return fmt.Sprintf("AS%s\\n%s", asn, strings.Join(records, " "))
+	}
+}
+
 func birdRouteToGraphviz(servers []string, responses []string, target string) string {
 	graph := make(map[string]string)
 	// Helper to add an edge
@@ -64,12 +74,12 @@ func birdRouteToGraphviz(servers []string, responses []string, target string) st
 					// Edge from originating server to nexthop
 					addEdge(server, "Nexthop:\\n"+routeNexthop, (map[bool]string{true: "[color=red]"})[routePreferred])
 					// and from nexthop to AS
-					addEdge("Nexthop:\\n"+routeNexthop, "AS"+paths[0], (map[bool]string{true: "[color=red]"})[routePreferred])
+					addEdge("Nexthop:\\n"+routeNexthop, getASNRepresentation(paths[0]), (map[bool]string{true: "[color=red]"})[routePreferred])
 					addPoint("Nexthop:\\n"+routeNexthop, "[shape=diamond]")
 					routeFound = true
 				} else {
 					// Edge from originating server to AS
-					addEdge(server, "AS"+paths[0], (map[bool]string{true: "[color=red]"})[routePreferred])
+					addEdge(server, getASNRepresentation(paths[0]), (map[bool]string{true: "[color=red]"})[routePreferred])
 					routeFound = true
 				}
 			}
@@ -79,10 +89,10 @@ func birdRouteToGraphviz(servers []string, responses []string, target string) st
 				if pathIndex == 0 {
 					continue
 				}
-				addEdge("AS"+paths[pathIndex-1], "AS"+paths[pathIndex], (map[bool]string{true: "[color=red]"})[routePreferred])
+				addEdge(getASNRepresentation(paths[pathIndex-1]), getASNRepresentation(paths[pathIndex]), (map[bool]string{true: "[color=red]"})[routePreferred])
 			}
 			// Last AS to destination
-			addEdge("AS"+paths[len(paths)-1], "Target: "+target, (map[bool]string{true: "[color=red]"})[routePreferred])
+			addEdge(getASNRepresentation(paths[len(paths)-1]), "Target: "+target, (map[bool]string{true: "[color=red]"})[routePreferred])
 		}
 
 		if !routeFound {
diff --git a/frontend/main.go b/frontend/main.go
index a793663..415af8c 100644
--- a/frontend/main.go
+++ b/frontend/main.go
@@ -8,37 +8,46 @@ import (
 )
 
 type settingType struct {
-	servers     []string
-	domain      string
-	proxyPort   int
-	whoisServer string
-	listen      string
+	servers      []string
+	domain       string
+	proxyPort    int
+	whoisServer  string
+	listen       string
+	dnsInterface string
 }
 
 var setting settingType
 
 func main() {
 	var settingDefault = settingType{
-		[]string{""}, "", 8000, "whois.verisign-grs.com", ":5000",
+		[]string{""},
+		"",
+		8000,
+		"whois.verisign-grs.com",
+		":5000",
+		"asn.cymru.com",
 	}
 
-	if serversEnv := os.Getenv("BIRDLG_SERVERS"); serversEnv != "" {
-		settingDefault.servers = strings.Split(serversEnv, ",")
+	if env := os.Getenv("BIRDLG_SERVERS"); env != "" {
+		settingDefault.servers = strings.Split(env, ",")
 	}
-	if domainEnv := os.Getenv("BIRDLG_DOMAIN"); domainEnv != "" {
-		settingDefault.domain = domainEnv
+	if env := os.Getenv("BIRDLG_DOMAIN"); env != "" {
+		settingDefault.domain = env
 	}
-	if proxyPortEnv := os.Getenv("BIRDLG_PROXY_PORT"); proxyPortEnv != "" {
+	if env := os.Getenv("BIRDLG_PROXY_PORT"); env != "" {
 		var err error
-		if settingDefault.proxyPort, err = strconv.Atoi(proxyPortEnv); err != nil {
+		if settingDefault.proxyPort, err = strconv.Atoi(env); err != nil {
 			panic(err)
 		}
 	}
-	if whoisEnv := os.Getenv("BIRDLG_WHOIS"); whoisEnv != "" {
-		settingDefault.whoisServer = whoisEnv
+	if env := os.Getenv("BIRDLG_WHOIS"); env != "" {
+		settingDefault.whoisServer = env
 	}
-	if listenEnv := os.Getenv("BIRDLG_LISTEN"); listenEnv != "" {
-		settingDefault.listen = listenEnv
+	if env := os.Getenv("BIRDLG_LISTEN"); env != "" {
+		settingDefault.listen = env
+	}
+	if env := os.Getenv("BIRDLG_DNS_INTERFACE"); env != "" {
+		settingDefault.dnsInterface = env
 	}
 
 	serversPtr := flag.String("servers", strings.Join(settingDefault.servers, ","), "server name prefixes, separated by comma")
@@ -46,6 +55,7 @@ func main() {
 	proxyPortPtr := flag.Int("proxy-port", settingDefault.proxyPort, "port bird-lgproxy is running on")
 	whoisPtr := flag.String("whois", settingDefault.whoisServer, "whois server for queries")
 	listenPtr := flag.String("listen", settingDefault.listen, "address bird-lg is listening on")
+	dnsInterfacePtr := flag.String("dns-interface", settingDefault.dnsInterface, "dns zone to query ASN information")
 	flag.Parse()
 
 	if *serversPtr == "" {
@@ -60,6 +70,7 @@ func main() {
 		*proxyPortPtr,
 		*whoisPtr,
 		*listenPtr,
+		*dnsInterfacePtr,
 	}
 
 	webServerStart()