Implement new 'discover' command

api/driver.go

@@ -46,6 +46,11 @@ func CheckHost(uuid string) (hosts.Host, error) {
 	return api.CheckHost(uuid)
 }
 
+// ListHosts asks the server for all the enrolled osquery nodes
+func ListHosts() (utils.Rows, error) {
+	return api.ListHosts()
+}
+
 // ScheduleQuery posts a query for the target host that osquery will poll for
 func ScheduleQuery(uuid string, query string) (string, error) {
 	return api.ScheduleQuery(uuid, query)

api/mock/https.go

@@ -17,6 +17,7 @@ import (
 	"github.com/AbGuthrie/goquery/api/models"
 	"github.com/AbGuthrie/goquery/config"
 	"github.com/AbGuthrie/goquery/hosts"
+	"github.com/AbGuthrie/goquery/utils"
 
 	"golang.org/x/crypto/ssh/terminal"
 )
@@ -94,9 +95,6 @@ func extractSSOResponse(response *http.Response) (string, string, error) {
 	if strings.Index(bodyStr, "Invalid username or password") != -1 {
 		return "", "", fmt.Errorf("Credential Failure")
 	}
-	if config.GetDebug() {
-		fmt.Printf("ssoResponse: %s\n", bodyStr)
-	}
 	// Hacky Extracts
 	loc := strings.Index(bodyStr, "name=\"SAMLResponse\"")
 	endLoc := strings.Index(bodyStr[loc+27:], "\" ")
@@ -123,10 +121,6 @@ func (instance *mockAPI) authenticate() error {
 		return nil
 	}
 
-	if config.GetDebug() {
-		fmt.Printf("ssoRequest: %s\nrelayState: %s\n", ssoRequest, relayState)
-	}
-
 	username, password := credentials()
 
 	response, err = instance.Client.PostForm("http://localhost:8002/sso",
@@ -137,10 +131,6 @@ func (instance *mockAPI) authenticate() error {
 	}
 
 	samlResponse, relayState, err := extractSSOResponse(response)
-	if config.GetDebug() {
-		fmt.Printf("ssoResponse: %s\nrelayState: %s\n", samlResponse, relayState)
-	}
-
 	if err != nil {
 		return err
 	}
@@ -215,6 +205,44 @@ func (instance *mockAPI) CheckHost(uuid string) (hosts.Host, error) {
 	}, nil
 }
 
+func (instance *mockAPI) ListHosts() (utils.Rows, error) {
+	if !instance.Authed {
+		err := instance.authenticate()
+		if err != nil {
+			return utils.Rows{}, err
+		}
+	}
+
+	response, err := instance.Client.Get("https://localhost:8001/listHosts")
+	if err != nil {
+		// Possible Authentication Failure
+		return utils.Rows{}, fmt.Errorf("ListHosts call failed: %s", err)
+	}
+	if response.StatusCode == 404 {
+		return utils.Rows{}, fmt.Errorf("Unknown Host")
+	}
+	if response.StatusCode != 200 {
+		return utils.Rows{}, fmt.Errorf("Server returned unknown error: %d", response.StatusCode)
+	}
+
+	bodyBytes, err := ioutil.ReadAll(response.Body)
+	if err != nil {
+		return utils.Rows{}, fmt.Errorf("Could not read response")
+	}
+	var hostResponse utils.Rows
+	err = json.Unmarshal(bodyBytes, &hostResponse)
+	if err != nil {
+		if config.GetDebug() {
+			fmt.Printf("Returned Body: %s\n", string(bodyBytes))
+		}
+		// Probable authentication failure
+		instance.Authed = false
+		return utils.Rows{}, err
+	}
+
+	return hostResponse, nil
+}
+
 func (instance *mockAPI) ScheduleQuery(uuid string, query string) (string, error) {
 	if !instance.Authed {
 		err := instance.authenticate()

api/models/models.go

@@ -10,6 +10,7 @@ import (
 // is blind to the implementation for code separation purposes.
 type GoQueryAPI interface {
 	CheckHost(string) (hosts.Host, error)
+	ListHosts() (utils.Rows, error)
 	ScheduleQuery(string, string) (string, error)
 	FetchResults(string) (utils.Rows, string, error)
 }

api/osctrl/https.go

@@ -16,6 +16,7 @@ import (
 
 	"github.com/AbGuthrie/goquery/api/models"
 	"github.com/AbGuthrie/goquery/config"
+	"github.com/AbGuthrie/goquery/utils"
 
 	"github.com/AbGuthrie/goquery/hosts"
 
@@ -190,6 +191,11 @@ func (instance *osctrlAPI) CheckHost(uuid string) (hosts.Host, error) {
 	}, nil
 }
 
+func (instance *osctrlAPI) ListHosts() (utils.Rows, error) {
+	// TODO: integrate osctrl
+	return utils.Rows{}, nil
+}
+
 func (instance *osctrlAPI) ScheduleQuery(uuid string, query string) (string, error) {
 	if !instance.Authed {
 		err := instance.authenticate()

cmd/main.go

@@ -124,9 +124,9 @@ func completer(in prompt.Document) []prompt.Suggest {
 				if len(description) == 0 {
 					description = alias.Command
 				}
-				prompts = append(prompts, prompt.Suggest{suggestion, description})
+				prompts = append(prompts, prompt.Suggest{Text: suggestion, Description: alias.Command})
 			} else if command, ok := commands.CommandMap[suggestion]; ok {
-				prompts = append(prompts, prompt.Suggest{suggestion, command.Help()})
+				prompts = append(prompts, prompt.Suggest{Text: suggestion, Description: command.Help()})
 			}
 		}
 		return prompt.FilterHasPrefix(prompts, command, true)

commands/command_map.go

@@ -27,6 +27,7 @@ func init() {
 		".connect":    GoQueryCommand{connect, connectHelp, connectSuggest},
 		".clear":      GoQueryCommand{clear, clearHelp, clearSuggest},
 		".disconnect": GoQueryCommand{disconnect, disconnectHelp, disconnectSuggest},
+		".discover":   GoQueryCommand{discover, discoverHelp, discoverSuggest},
 		".exit":       GoQueryCommand{exit, exitHelp, exitSuggest},
 		".help":       GoQueryCommand{help, helpHelp, helpSuggest},
 		".history":    GoQueryCommand{history, historyHelp, historySuggest},

commands/connect.go

@@ -55,7 +55,10 @@ func connectHelp() string {
 func connectSuggest(cmdline string) []prompt.Suggest {
 	prompts := []prompt.Suggest{}
 	for _, host := range hosts.GetCurrentHosts() {
-		prompts = append(prompts, prompt.Suggest{host.UUID, host.ComputerName})
+		prompts = append(prompts, prompt.Suggest{
+			Text:        host.UUID,
+			Description: host.ComputerName,
+		})
 	}
 	return prompts
 }

commands/disconnect.go

@@ -31,7 +31,10 @@ func disconnectHelp() string {
 func disconnectSuggest(cmdline string) []prompt.Suggest {
 	prompts := []prompt.Suggest{}
 	for _, host := range hosts.GetCurrentHosts() {
-		prompts = append(prompts, prompt.Suggest{host.UUID, host.ComputerName})
+		prompts = append(prompts, prompt.Suggest{
+			Text:        host.UUID,
+			Description: host.ComputerName,
+		})
 	}
 	return prompts
 }

commands/discover.go

@@ -0,0 +1,36 @@
+package commands
+
+import (
+	"fmt"
+	"strings"
+
+	"github.com/AbGuthrie/goquery/api"
+	"github.com/AbGuthrie/goquery/utils"
+
+	prompt "github.com/c-bata/go-prompt"
+)
+
+func discover(cmdline string) error {
+	args := strings.Split(cmdline, " ") // Separate command and arguments
+	if len(args) != 1 {
+		return fmt.Errorf("This command takes no parameters")
+	}
+
+	// Query for list of available hosts
+	enrolledHosts, err := api.ListHosts()
+	if err != nil {
+		return fmt.Errorf("Error querying available hosts: %s", err)
+	}
+
+	utils.PrettyPrintQueryResults(enrolledHosts)
+
+	return nil
+}
+
+func discoverHelp() string {
+	return "Prints all hosts registered with api server"
+}
+
+func discoverSuggest(cmdline string) []prompt.Suggest {
+	return []prompt.Suggest{}
+}

commands/mode.go

@@ -54,7 +54,10 @@ func changeModeSuggest(cmdline string) []prompt.Suggest {
 	sort.Strings(modeNames)
 
 	for _, mode := range modeNames {
-		prompts = append(prompts, prompt.Suggest{mode, ""})
+		prompts = append(prompts, prompt.Suggest{
+			Text:        mode,
+			Description: "",
+		})
 	}
 
 	return prompts

commands/query.go

@@ -58,7 +58,10 @@ func querySuggest(cmdline string) []prompt.Suggest {
 		return prompts
 	}
 	for _, table := range host.Tables {
-		prompts = append(prompts, prompt.Suggest{table, ""})
+		prompts = append(prompts, prompt.Suggest{
+			Text:        table,
+			Description: "",
+		})
 	}
 
 	return prompts

commands/resume.go

@@ -45,7 +45,10 @@ func resumeSuggest(cmdline string) []prompt.Suggest {
 	}
 
 	for _, query := range host.QueryHistory {
-		prompts = append(prompts, prompt.Suggest{query.Name, query.SQL})
+		prompts = append(prompts, prompt.Suggest{
+			Text:        query.Name,
+			Description: query.SQL,
+		})
 	}
 	return prompts
 }

goserver/mock_osquery_server.go

@@ -14,6 +14,7 @@ import (
 	"math/rand"
 	"net/http"
 	"net/url"
+	"sort"
 	"strings"
 	"time"
 
@@ -290,6 +291,26 @@ func checkHost(w http.ResponseWriter, r *http.Request) {
 	fmt.Fprintf(w, "%s", renderedHost)
 }
 
+func listHosts(w http.ResponseWriter, r *http.Request) {
+	fmt.Printf("List hosts call")
+	hosts := []Host{}
+
+	for _, enrolledHost := range enrolledHosts {
+		hosts = append(hosts, enrolledHost)
+	}
+
+	sort.SliceStable(hosts, func(i, j int) bool {
+		return hosts[i].ComputerName > hosts[j].ComputerName
+	})
+
+	hostsJSON, err := json.Marshal(hosts)
+	if err != nil {
+		w.WriteHeader(http.StatusInternalServerError)
+		return
+	}
+	fmt.Fprintf(w, "%s", hostsJSON)
+}
+
 func scheduleQuery(w http.ResponseWriter, r *http.Request) {
 	uuid := r.FormValue("uuid")
 	sentQuery, err := json.Marshal(r.FormValue("query"))
@@ -416,16 +437,19 @@ func main() {
 		}
 		fmt.Printf("Registered ourselves with the IDP Service\n")
 
-		ch := http.HandlerFunc(checkHost)
-		sq := http.HandlerFunc(scheduleQuery)
-		fr := http.HandlerFunc(fetchResults)
+		_checkHost := http.HandlerFunc(checkHost)
+		_listHosts := http.HandlerFunc(listHosts)
+		_scheduleQuery := http.HandlerFunc(scheduleQuery)
+		_fetchResults := http.HandlerFunc(fetchResults)
 
-		http.Handle("/checkHost", samlSP.RequireAccount(ch))
-		http.Handle("/scheduleQuery", samlSP.RequireAccount(sq))
-		http.Handle("/fetchResults", samlSP.RequireAccount(fr))
+		http.Handle("/checkHost", samlSP.RequireAccount(_checkHost))
+		http.Handle("/listHosts", _listHosts)
+		http.Handle("/scheduleQuery", samlSP.RequireAccount(_scheduleQuery))
+		http.Handle("/fetchResults", samlSP.RequireAccount(_fetchResults))
 		http.Handle("/saml/", samlSP)
 	} else {
 		http.HandleFunc("/checkHost", checkHost)
+		http.HandleFunc("/listHosts", listHosts)
 		http.HandleFunc("/scheduleQuery", scheduleQuery)
 		http.HandleFunc("/fetchResults", fetchResults)
 	}