From a79af2ba3175f86e00989f5f31c0e4664b94f716 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?J=C3=B6rg=20Thalheim?= <joerg@higgsboson.tk>
Date: Sat, 23 Jan 2016 21:52:20 +0000
Subject: [PATCH] implement type filtering

---
 README.md      |  1 +
 whois/flags.go | 15 +++++++++++++++
 whois/query.go | 12 ++++++++----
 3 files changed, 24 insertions(+), 4 deletions(-)

diff --git a/README.md b/README.md
index 2ee21b9..634f3bc 100644
--- a/README.md
+++ b/README.md
@@ -61,6 +61,7 @@ it run as user nobody, who cannot bind to port 43 itself.
 - version: `$ whois -h <server> -q version`
 - sources: `$ whois -h <server> -q sources`
 - types: `$ whois -h <server> -q types`
+- type filtering: `$ whois -h 172.23.75.6 -T aut-num,person Mic92-DN42 AS4242420092`
 
 
 ## TODO
diff --git a/whois/flags.go b/whois/flags.go
index b912a28..1f8a13c 100644
--- a/whois/flags.go
+++ b/whois/flags.go
@@ -36,17 +36,32 @@ func shellSplit(s string) []string {
 
 type Flags struct {
 	ServerInfo string
+	TypeSchema string
+	Types      map[string]bool
 	Args       []string
 }
 
 func parseFlags(request string) (*Flags, *flag.FlagSet, error) {
 	args := shellSplit(request)
 	set := flag.NewFlagSet("whois42d", flag.ContinueOnError)
+	var typeField string
 	f := Flags{}
 	set.StringVar(&f.ServerInfo, "q", "", "[version|sources|types] query specified server info")
+	set.StringVar(&f.TypeSchema, "t", "", "request template for object of TYPE")
+	set.StringVar(&typeField, "T", "", "TYPE[,TYPE]... only look for objects of TYPE")
+
 	if err := set.Parse(args); err != nil {
 		return nil, set, err
 	}
+
+	if typeField != "" {
+		types := strings.Split(typeField, ",")
+		f.Types = make(map[string]bool, len(types))
+		for _, t := range types {
+			f.Types[t] = true
+		}
+	}
+
 	f.Args = set.Args()
 	return &f, set, nil
 }
diff --git a/whois/query.go b/whois/query.go
index ff5ae12..3a783b3 100644
--- a/whois/query.go
+++ b/whois/query.go
@@ -56,9 +56,13 @@ var whoisTypes = []Type{
 	{"as-block", regexp.MustCompile(`\d+_\d+`), UPPER},
 }
 
-func (r *Registry) handleObject(conn *net.TCPConn, object Object) bool {
+func (r *Registry) handleObject(conn *net.TCPConn, object Object, flags *Flags) bool {
 	found := false
 	for _, t := range whoisTypes {
+		if len(flags.Types) > 0 && !flags.Types[t.Name] {
+			continue
+		}
+
 		if t.Kind == ROUTE || t.Kind == ROUTE6 {
 			if object[t.Kind] != nil {
 				found = found || r.printNet(conn, t.Name, object[t.Kind].(net.IP))
@@ -84,12 +88,12 @@ func (r *Registry) HandleQuery(conn *net.TCPConn) {
 
 	flags := query.Flags
 	if flags.ServerInfo != "" {
-		r.printServerInfo(conn, strings.TrimSpace(flags.ServerInfo))
+		printServerInfo(conn, strings.TrimSpace(flags.ServerInfo))
 		return
 	}
 	found := false
 	for _, obj := range query.Objects {
-		if r.handleObject(conn, obj) {
+		if r.handleObject(conn, obj, flags) {
 			found = true
 		}
 	}
@@ -176,7 +180,7 @@ func parseQuery(conn *net.TCPConn) *Query {
 	return &query
 }
 
-func (r *Registry) printServerInfo(conn *net.TCPConn, what string) {
+func printServerInfo(conn *net.TCPConn, what string) {
 	switch what {
 	case "version":
 		fmt.Fprintf(conn, "%% whois42d v%d\n", VERSION)