Added TCP protocol test.

2 files changed, 257 insertions, 0 deletions

diff --git a/lib/custodian/protocol-tests/tcp.rb b/lib/custodian/protocol-tests/tcp.rb
new file mode 100755
index 0000000..847062c
--- /dev/null
+++ b/lib/custodian/protocol-tests/tcp.rb
@@ -0,0 +1,163 @@
+#!/usr/bin/ruby1.8
+
+
+
+require 'socket'
+require 'timeout'
+
+
+#
+# Test that we can receive a response from a TCP server that matches
+# a given banner.
+#
+class TCPTest
+
+  #
+  # Data passed from the JSON hash.
+  #
+  attr_reader :test_data
+
+  #
+  # The error text we return on failure.
+  #
+  attr_reader :error
+
+
+
+  #
+  # Save the data away.
+  #
+  def initialize( data )
+    @test_data = data
+    @error     = nil
+
+
+    #
+    # Ensure we have a host to probe
+    #
+    if ( @test_data["target_host"].nil? )
+      @error = "Missing target for the test."
+      raise ArgumentError, @error
+    end
+
+    #
+    # Ensure we have a port to test.
+    #
+    if ( @test_data["test_port"].nil? )
+      @error = "Missing port for the test."
+      raise ArgumentError, @error
+    end
+  end
+
+
+  #
+  # Run the test.
+  #
+  #  Return "true" on success
+  #
+  #  Return "false" on failure.
+  #
+  # If the test fails the details should be retrieved from "error".
+  #
+  def run_test
+    #
+    # Reset state from previous test.
+    #
+    @error = nil
+
+    #
+    #  Get the hostname & port to test against.
+    #
+    host = @test_data['target_host']
+    port = @test_data['test_port']
+
+    #
+    #  Get the banner we expect
+    #
+    banner = @test_data['banner']
+
+    puts "TCP testing host #{host}:#{port}" if ( @test_data['verbose'] )
+    if ( @test_data['verbose'] && ( !banner.nil? ) )
+        puts "Looking for banner '#{banner}'"
+    end
+
+    begin
+      timeout(@test_data["timeout"].to_i) do
+
+        begin
+          socket = TCPSocket.new( host, port )
+          socket.puts( "QUIT")
+
+          # read a banner from the remote server
+          read = socket.gets(nil)
+          read = read[0,255]
+
+          socket.close()
+
+          if ( !banner.nil? )
+              # test for bannerbanner
+              if ( read =~ /#{banner}/i )
+                   puts "We connected and matched the banner against '#{read}'" if ( @test_data['verbose'] )
+                return true
+              end
+
+              @error = "We expected a banner matching '#{banner}' but we got '#{read}'"
+              return false
+          end
+       rescue
+          @error = "Exception connecting to host #{host}:#{port} - #{$!}"
+          return false
+        end
+      end
+    rescue Timeout::Error => e
+      @error = "TIMEOUT: #{e}"
+      return false
+    end
+
+    @error = "Misc failure"
+    return false
+  end
+
+
+
+  #
+  #  Return the error text for why this test failed.
+  #
+  def error
+    return @error
+  end
+
+end
+
+
+#
+# Sample test, for testing.
+#
+if __FILE__ == $0 then
+
+  #
+  #  Sample data.
+  #
+  test = {
+    "target_host" => "mail.steve.org.uk",
+    "test_type"   => "tcp",
+    "test_port"   => "25",
+    "banner"      => "SMTP",
+    "verbose"     => 1,
+    "timeout"     => 5,
+    "test_alert"  => "TCP-port failure",
+  }
+
+
+  #
+  #  Run the test.
+  #
+  obj = TCPTest.new( test )
+  if ( obj.run_test )
+    puts "TEST OK"
+  else
+    puts "TEST FAILED"
+    puts obj.error()
+  end
+
+end
diff --git a/t/test-protocol-tcp.rb b/t/test-protocol-tcp.rb
new file mode 100755
index 0000000..655b6a1
--- /dev/null
+++ b/t/test-protocol-tcp.rb
@@ -0,0 +1,94 @@
+#!/usr/bin/ruby -I../lib/ -Ilib/
+
+
+require 'test/unit'
+require 'custodian/protocol-tests/tcp.rb'
+
+
+
+
+#
+# Unit test for the TCP-protocol probe.
+#
+class TestTCPProtocolProbe < Test::Unit::TestCase
+
+  #
+  # Create the test suite environment: NOP.
+  #
+  def setup
+  end
+
+  #
+  # Destroy the test suite environment: NOP.
+  #
+  def teardown
+  end
+
+
+  #
+  #  Test we can create a new TCPTest object.
+  #
+  def test_init
+    test_data_good = {
+      "target_host" => "smtp.bytemark.co.uk",
+      "test_type"   => "tcp",
+      "banner"      => 220,
+      "verbose"     => 1,
+      "test_port"   => 25,
+      "test_alert"  => "SMTP service down",
+    }
+
+    #
+    # Missing a port number
+    #
+    test_data_bad_one = {
+      "target_host" => "smtp.bytemark.co.uk",
+      "test_type"   => "tcp",
+      "banner"      => 220,
+      "verbose"     => 1,
+      "test_alert"  => "SMTP service down",
+    }
+
+    #
+    #  Missing a target host to probe
+    #
+    test_data_bad_two = {
+      "test_type"   => "tcp",
+      "banner"      => 220,
+      "verbose"     => 1,
+      "test_port"   => 25,
+      "test_alert"  => "SMTP service down",
+    }
+
+
+    #
+    #  Create a new FTPTest object.  This should succeed
+    #
+    good = TCPTest.new( test_data_good )
+    assert( good )
+
+    #
+    #  There will be no error setup
+    #
+    assert( good.error().nil? )
+
+    #
+    #  Now create a probe with a missing port.
+    #
+    assert_raise ArgumentError do
+      bad = TCPTest.new( test_data_bad_one )
+    end
+
+
+    #
+    #  Now create a probe with a missing host.
+    #
+    assert_raise ArgumentError do
+      bad = TCPTest.new( test_data_bad_two )
+    end
+
+  end
+
+
+
+end