fw_lab_triage_helper: support arbitrary pools

BUG=None
TEST=Ran with default pool, and with a pool (-pool faft-test) known
to contain non chromeos1 DUTs (so as to exercise the filtering).

Change-Id: I3555e93f0c8a78ecfcb9954e77aba4681ae020a6
Reviewed-on: https://chromium-review.googlesource.com/1741611
Tested-by: Kevin Shelton <kmshelton@chromium.org>
Commit-Ready: Kevin Shelton <kmshelton@chromium.org>
Legacy-Commit-Queue: Commit Bot <commit-bot@chromium.org>
Reviewed-by: Greg Edelston <gredelston@google.com>
diff --git a/provingground/firmware/fw_lab_triage_helper.go b/provingground/firmware/fw_lab_triage_helper.go
index 791d8aa..de010e9 100644
--- a/provingground/firmware/fw_lab_triage_helper.go
+++ b/provingground/firmware/fw_lab_triage_helper.go
@@ -5,6 +5,7 @@
 package main
 
 import (
+	"flag"
 	"fmt"
 	"log"
 	"os"
@@ -76,22 +77,24 @@
 func main() {
 	duts := []dut{}
 
+	poolPtr := flag.String("pool", "faft-cr50", "A pool of DUTs to operate on.")
+	flag.Parse()
+
 	fmt.Println(warningMessage)
 
 	log.Print("Gathering DUT info via atest...")
-	// TODO(kmshelton): Support arbitrary pools.  Remember to sanitize for chromeos1 (high
-	// touch lab) when adding arbitrary pool support, as we would not want to do operations
-	// like ssh'ing to labstations in the low touch lab.  The current hardcoded pool, it is
-	// safe to assume no low touch lab devices are operated upon.
-	cmd := exec.Command("atest", "host", "list", "--hostnames-only", "--label=pool:faft-cr50")
+
+	cmd := exec.Command("atest", "host", "list", "--hostnames-only", "--label=pool:"+*poolPtr)
 	out, err := cmd.Output()
 	if err != nil {
 		log.Fatalf("<atest host list> encountered: %s", err)
 	}
-	hostnames := strings.Fields(string(out))
 
-	for _, hostname := range hostnames {
-		duts = append(duts, newDut(hostname))
+	// Removing hostnames that don't begin with "chromeos1-" removes those that are not in the firmware lab.
+	for _, hostname := range strings.Fields(string(out)) {
+		if strings.HasPrefix(hostname, "chromeos1-") {
+			duts = append(duts, newDut(hostname))
+		}
 	}
 
 	log.Print("Summarizing DUT info...")