From 5226cd0c616094e548b8eee0bf581f6c17eb70d8 Mon Sep 17 00:00:00 2001
From: Zach <zgiles@gmail.com>
Date: Sun, 2 Sep 2018 17:30:46 -0400
Subject: [PATCH] Generalize object retrieval

Generalize object retrieval in to objects and string pairs to enable multiple retrieval patterns ( like HTTP )
---
 whois/query.go    | 56 +++++++++++++++++++++++++++++++++--------------
 whois/registry.go | 14 ++++++++++++
 2 files changed, 53 insertions(+), 17 deletions(-)

diff --git a/whois/query.go b/whois/query.go
index 855c00a..6fc6483 100644
--- a/whois/query.go
+++ b/whois/query.go
@@ -10,26 +10,54 @@ import (
 	"strings"
 )
 
+type pathpair struct {
+	objtype string
+	obj		string
+}
+
 func (r *Registry) handleObject(conn *net.TCPConn, object Object, flags *Flags) bool {
 	found := false
-	for _, t := range r.whoisTypes {
-		if len(flags.Types) > 0 && !flags.Types[t.Name] {
-			continue
-		}
+	paths := r.findObjectPaths(object)
+	for _, p := range paths {
+			r.printObject(conn, p.objtype, p.obj)
+			found = true
+	}
+	return found
+}
 
+func (r *Registry) findObjectPaths(object Object) []pathpair {
+	var paths []pathpair
+	for _, t := range r.whoisTypes {
 		if t.Kind == ROUTE || t.Kind == ROUTE6 {
 			if object[t.Kind] != nil {
-				found = found || r.printNet(conn, t.Name, object[t.Kind].(net.IP))
+				p := r.getObjFromIP(t.Name, object[t.Kind].(net.IP))
+				paths = append(paths, p...)
 			}
 		} else {
 			arg := object[t.Kind].(string)
 			if t.Pattern.MatchString(arg) {
-				r.printObject(conn, t.Name, arg)
-				found = true
+				paths = append(paths, pathpair{t.Name, arg})
 			}
 		}
 	}
-	return found
+	return paths
+}
+
+func (r *Registry) getObjFromIP(objType string, ip net.IP) []pathpair {
+	var paths []pathpair
+	routePath := path.Join(r.DataPath, objType)
+	cidrs, err := readCidrs(routePath)
+	if err != nil {
+		return paths
+	}
+
+	for _, c := range cidrs {
+		if c.Contains(ip) {
+			obj := strings.Replace(c.String(), "/", "_", -1)
+			paths = append(paths, pathpair{objType, obj})
+		}
+	}
+	return paths
 }
 
 func (r *Registry) HandleQuery(conn *net.TCPConn) {
@@ -119,18 +147,12 @@ func (r *Registry) printNet(conn *net.TCPConn, name string, ip net.IP) bool {
 }
 
 func (r *Registry) printObject(conn *net.TCPConn, objType string, obj string) {
-	file := path.Join(r.DataPath, objType, obj)
-
-	f, err := os.Open(file)
-	defer f.Close()
+	content, path, err := r.retrieveObject(objType, obj)
 	if err != nil {
-		if os.IsNotExist(err) {
-			return
-		}
 		fmt.Fprintf(os.Stderr, "Error: %s\n", err)
 		return
 	}
-	fmt.Fprintf(conn, "%% Information related to '%s':\n", file[len(r.DataPath)+1:])
-	conn.ReadFrom(f)
+	fmt.Fprintf(conn, "%% Information related to '%s':\n", path)
+	conn.Write(content)
 	fmt.Fprint(conn, "\n")
 }
diff --git a/whois/registry.go b/whois/registry.go
index be2cb0f..f5c12d0 100644
--- a/whois/registry.go
+++ b/whois/registry.go
@@ -62,6 +62,20 @@ func New(DataPath string, Header string, DNSTopLevel string, RegistryTopLevel st
 	return r
 }
 
+func (r *Registry) retrieveObject(objType string, obj string) ([]byte, string, error) {
+	file := path.Join(r.DataPath, objType, obj)
+	f, err := os.Open(file)
+	defer f.Close()
+	if err != nil {
+		return []byte(""), string(""), err
+	}
+	fall, fallerr := ioutil.ReadAll(f)
+	if fallerr != nil {
+		return []byte(""), string(""), err
+	}
+	return fall, file[len(r.DataPath)+1:], nil
+}
+
 func readCidrs(path string) ([]net.IPNet, error) {
 	files, err := ioutil.ReadDir(path)
 	if err != nil {