Implement query_response matching for the tcp prober.

fixes #13

diff --git a/README.md b/README.md
index 8f1178a936a7611f0311377b5bfc2cbf68a49f2a..7b698f08b145a8aba2c11a77bd81e7d3c8ed6382 100644 (file)
--- a/README.md
+++ b/README.md
@@ -39,6 +39,22 @@ modules:
   tcp_connect:
     prober: tcp
     timeout: 5s
+  ssh_banner:
+    prober: tcp
+    timeout: 5s
+    tcp:
+      query_response:
+      - expect: "^SSH-2.0-"
+  irc_banner:
+    prober: tcp
+    timeout: 5s
+    tcp:
+      query_response:
+      - send: "NICK prober"
+      - send: "USER prober prober prober :prober"
+      - expect: "PING :([^ ]+)"
+        send: "PONG ${1}"
+      - expect: "^:[^ ]+ 001"
   icmp:
     prober: icmp
     timeout: 5s

diff --git a/main.go b/main.go
index f31a539e8dec958737f8e50bf29394b7838903ef..b5ae316582fbf2bf9a724f46c643252432a0b5c8 100644 (file)
--- a/main.go
+++ b/main.go
@@ -39,7 +39,13 @@ type HTTPProbe struct {
        FailIfNotMatchesRegexp []string `yaml:"fail_if_not_matches_regexp"`
 }
 
+type QueryResponse struct {
+       Expect string `yaml:"expect"`
+       Send   string `yaml:"send"`
+}
+
 type TCPProbe struct {
+       QueryResponse []QueryResponse `yaml:"query_response"`
 }
 
 type ICMPProbe struct {

diff --git a/tcp.go b/tcp.go
index 685de9e00c481a395ad1220083ab9a630b5fc961..ff8950dee201183503381e005b639236106065df 100644 (file)
--- a/tcp.go
+++ b/tcp.go
@@ -1,15 +1,63 @@
 package main
 
 import (
+       "bufio"
+       "fmt"
        "net"
        "net/http"
+       "regexp"
+       "time"
+
+       "github.com/prometheus/log"
 )
 
-func probeTCP(target string, w http.ResponseWriter, module Module) (success bool) {
+func probeTCP(target string, w http.ResponseWriter, module Module) bool {
+       deadline := time.Now().Add(module.Timeout)
        conn, err := net.DialTimeout("tcp", target, module.Timeout)
-       if err == nil {
-               success = true
-               conn.Close()
+       if err != nil {
+               return false
+       }
+       defer conn.Close()
+       // Set a deadline to prevent the following code from blocking forever.
+       // If a deadline cannot be set, better fail the probe by returning an error
+       // now rather than blocking forever.
+       if err := conn.SetDeadline(deadline); err != nil {
+               return false
+       }
+       scanner := bufio.NewScanner(conn)
+       for _, qr := range module.TCP.QueryResponse {
+               log.Debugf("Processing query response entry %+v", qr)
+               send := qr.Send
+               if qr.Expect != "" {
+                       re, err := regexp.Compile(qr.Expect)
+                       if err != nil {
+                               log.Errorf("Could not compile %q into regular expression: %v", qr.Expect, err)
+                               return false
+                       }
+                       var match []int
+                       // Read lines until one of them matches the configured regexp.
+                       for scanner.Scan() {
+                               log.Debugf("read %q\n", scanner.Text())
+                               match = re.FindSubmatchIndex(scanner.Bytes())
+                               if match != nil {
+                                       log.Debugf("regexp %q matched %q", re, scanner.Text())
+                                       break
+                               }
+                       }
+                       if scanner.Err() != nil {
+                               return false
+                       }
+                       if match == nil {
+                               return false
+                       }
+                       send = string(re.Expand(nil, []byte(send), scanner.Bytes(), match))
+               }
+               if send != "" {
+                       log.Debugf("Sending %q", send)
+                       if _, err := fmt.Fprintf(conn, "%s\n", send); err != nil {
+                               return false
+                       }
+               }
        }
-       return
+       return true
 }

diff --git a/tcp_test.go b/tcp_test.go
index 9243e71673cd6185ec55c3b6cd4d30710a811160..70474e917638514d49ec79e249882d9ce36f9044 100644 (file)
--- a/tcp_test.go
+++ b/tcp_test.go
@@ -14,6 +14,7 @@
 package main
 
 import (
+       "fmt"
        "net"
        "testing"
        "time"
@@ -47,3 +48,99 @@ func TestTCPConnectionFails(t *testing.T) {
                t.Fatalf("TCP module suceeded, expected failure.")
        }
 }
+
+func TestTCPConnectionQueryResponseIRC(t *testing.T) {
+       ln, err := net.Listen("tcp", "localhost:0")
+       if err != nil {
+               t.Fatalf("Error listening on socket: %s", err)
+       }
+       defer ln.Close()
+
+       module := Module{
+               Timeout: time.Second,
+               TCP: TCPProbe{
+                       QueryResponse: []QueryResponse{
+                               QueryResponse{Send: "NICK prober"},
+                               QueryResponse{Send: "USER prober prober prober :prober"},
+                               QueryResponse{Expect: "^:[^ ]+ 001"},
+                       },
+               },
+       }
+
+       ch := make(chan (struct{}))
+       go func() {
+               conn, err := ln.Accept()
+               if err != nil {
+                       t.Fatalf("Error accepting on socket: %s", err)
+               }
+               fmt.Fprintf(conn, ":ircd.localhost NOTICE AUTH :*** Looking up your hostname...\n")
+               var nick, user, mode, unused, realname string
+               fmt.Fscanf(conn, "NICK %s", &nick)
+               fmt.Fscanf(conn, "USER %s %s %s :%s", &user, &mode, &unused, &realname)
+               fmt.Fprintf(conn, ":ircd.localhost 001 %s :Welcome to IRC!\n", nick)
+               conn.Close()
+               ch <- struct{}{}
+       }()
+       if !probeTCP(ln.Addr().String(), nil, module) {
+               t.Fatalf("TCP module failed, expected success.")
+       }
+       <-ch
+
+       go func() {
+               conn, err := ln.Accept()
+               if err != nil {
+                       t.Fatalf("Error accepting on socket: %s", err)
+               }
+               fmt.Fprintf(conn, ":ircd.localhost NOTICE AUTH :*** Looking up your hostname...\n")
+               var nick, user, mode, unused, realname string
+               fmt.Fscanf(conn, "NICK %s", &nick)
+               fmt.Fscanf(conn, "USER %s %s %s :%s", &user, &mode, &unused, &realname)
+               fmt.Fprintf(conn, "ERROR: Your IP address has been blacklisted.\n")
+               conn.Close()
+               ch <- struct{}{}
+       }()
+       if probeTCP(ln.Addr().String(), nil, module) {
+               t.Fatalf("TCP module succeeded, expected failure.")
+       }
+       <-ch
+}
+
+func TestTCPConnectionQueryResponseMatching(t *testing.T) {
+       ln, err := net.Listen("tcp", "localhost:0")
+       if err != nil {
+               t.Fatalf("Error listening on socket: %s", err)
+       }
+       defer ln.Close()
+
+       module := Module{
+               Timeout: time.Second,
+               TCP: TCPProbe{
+                       QueryResponse: []QueryResponse{
+                               QueryResponse{
+                                       Expect: "SSH-2.0-(OpenSSH_6.9p1) Debian-2",
+                                       Send:   "CONFIRM ${1}",
+                               },
+                       },
+               },
+       }
+
+       ch := make(chan string)
+       go func() {
+               conn, err := ln.Accept()
+               if err != nil {
+                       t.Fatalf("Error accepting on socket: %s", err)
+               }
+               conn.SetDeadline(time.Now().Add(1 * time.Second))
+               fmt.Fprintf(conn, "SSH-2.0-OpenSSH_6.9p1 Debian-2\n")
+               var version string
+               fmt.Fscanf(conn, "CONFIRM %s", &version)
+               conn.Close()
+               ch <- version
+       }()
+       if !probeTCP(ln.Addr().String(), nil, module) {
+               t.Fatalf("TCP module failed, expected success.")
+       }
+       if got, want := <-ch, "OpenSSH_6.9p1"; got != want {
+               t.Fatalf("Read unexpected version: got %q, want %q", got, want)
+       }
+}