From d2dc120e73f7d25f1ec0d8b58914f1f4ea6f5fd8 Mon Sep 17 00:00:00 2001 From: Michael Stapelberg Date: Tue, 1 Dec 2015 23:34:24 +0100 Subject: [PATCH] Implement query_response matching for the tcp prober. fixes #13 --- README.md | 16 +++++++++ main.go | 6 ++++ tcp.go | 58 +++++++++++++++++++++++++++++--- tcp_test.go | 97 +++++++++++++++++++++++++++++++++++++++++++++++++++++ 4 files changed, 172 insertions(+), 5 deletions(-) diff --git a/README.md b/README.md index 8f1178a..7b698f0 100644 --- 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 f31a539..b5ae316 100644 --- 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 685de9e..ff8950d 100644 --- 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 9243e71..70474e9 100644 --- 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) + } +} -- 2.25.1