diff options
| -rwxr-xr-x | lib/custodian/protocol-tests/tcp.rb | 163 | ||||
| -rwxr-xr-x | t/test-protocol-tcp.rb | 94 |
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 |