2 files changed, 180 insertions, 3 deletions
diff --git a/lib/custodian/protocoltest/ssh.rb b/lib/custodian/protocoltest/ssh.rb
new file mode 100644
index 0000000..73f23da
--- /dev/null
+++ b/lib/custodian/protocoltest/ssh.rb
@@ -0,0 +1,78 @@
+#
+#  The SSH-protocol test.
+#
+#  This object is instantiated if the parser sees a line such as:
+#
+###
+### foo.vm.bytemark.co.uk must run ssh on 22 otherwise 'ssh fail'.
+###
+#
+#  The specification of the port is mandatory.
+#
+class SSHTest < TCPTest
+
+  #
+  # The host to test against.
+  #
+  attr_reader :host
+
+
+  #
+  # The port to connect to.
+  #
+  attr_reader :port
+
+
+
+  #
+  # Constructor
+  #
+  # Ensure we received a port to run the test against.
+  #
+  def initialize( line )
+    #
+    # Save the host
+    #
+    @host  = line.split( /\s+/)[0]
+
+    #
+    # Save the port
+    #
+    if ( line =~ /on\s+([0-9]+)/ )
+      @port = $1.dup
+    else
+      @port = 22
+    end
+  end
+
+
+  #
+  # Helper for development.
+  #
+  def to_s
+    "ssh-test of #{@host}:#{@port}."
+  end
+
+
+  #
+  # Run the TCP-protocol test.
+  #
+  def run_test
+
+    # reset the error, in case we were previously executed.
+    @error = nil
+
+    run_test_internal( @host, @port, "SSH" )
+  end
+
+
+  #
+  # If the test fails then report the error.
+  #
+  def error
+    @error
+  end
+
+  register_test_type "ssh"
+
+end
diff --git a/lib/custodian/protocoltest/tcp.rb b/lib/custodian/protocoltest/tcp.rb
index fef6dea..115bf6a 100644
--- a/lib/custodian/protocoltest/tcp.rb
+++ b/lib/custodian/protocoltest/tcp.rb
@@ -1,3 +1,7 @@
+require 'socket'
+require 'timeout'
+
+
 #
 #  The TCP-protocol test.
 #
@@ -13,21 +17,70 @@ class TCPTest < ProtocolTest
 
 
   #
+  # The host to test against.
+  #
+  attr_reader :host
+
+
+  #
+  # The port to connect to.
+  #
+  attr_reader :port
+
+
+  #
+  #  The banner to look for, may be nil.
+  #
+  attr_reader :banner
+
+
+
+
+  #
   # Constructor
   #
   # Ensure we received a port to run the TCP-test against.
   #
   def initialize( line )
-    raise ArgumentError, "Missing port" unless ( line =~ /on\s+([0-9]+)/ );
+
+    #
+    # Save the host
+    #
+    @host  = line.split( /\s+/)[0]
+
+    #
+    # Save the port
+    #
+    if ( line =~ /on\s+([0-9]+)/ )
+      @port = $1.dup
+    else
+      @port = nil
+    end
+
+    #
+    # Save the optional banner.
+    #
+    if ( line =~ /with\s+banner\s+'([^']+)'/ )
+      @banner = $1.dup
+    else
+      @banner = nil
+    end
+
     @error = nil
+
+    if ( @port.nil? )
+      raise ArgumentError, "Missing port to test against"
+    end
   end
 
 
+
+
   #
   # Helper for development.
   #
-  def display
-    puts "I'm a TCP-test!"
+  def to_s
+    "tcp-test of #{@host}:#{@port} looking for banner '#{@banner}'."
   end
 
 
@@ -39,6 +92,52 @@ class TCPTest < ProtocolTest
     # reset the error, in case we were previously executed.
     @error = nil
 
+    return( run_test_internal( @host, @port, @banner ) )
+  end
+
+
+  #
+  #
+  #
+  def run_test_internal( host, port, banner )
+    begin
+      timeout(30) do
+        begin
+          socket = TCPSocket.new( host, port )
+          socket.puts( "QUIT")
+
+          # read a banner from the remote server
+          read = socket.gets(nil)
+
+          # trim to a sane length & strip newlines.
+          read = read[0,255] unless ( read.nil? )
+          read.gsub!(/[\n\r]/, "") unless ( read.nil? )
+
+          socket.close()
+
+          if ( banner.nil? )
+            @error = nil
+            return true
+          else
+            # test for banner
+            if ( ( !read.nil? ) && ( read =~ /#{banner}/i ) )
+              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