A HTTP(S) healthcheck script

#!/usr/bin/ruby

=begin
Purpose:
    Performs HTTP/HTTPs request to check whether some node responds
    in given time-frame. You can specify alternate connect IP/port
    combination (which is useful to test LVS clusters with VHOST
    based webservice) and also whether you want document body as
    output or its (sha1 or md5) digest.

    In current version also supports request timing (benchmark)
    and multiple urls.

Author: Wejn <wejn at box dot cz>
License: GPLv2 (without the "latter" option)
Requires: Ruby >= 1.8
TS: 20080403013000

Background info: https://wejn/org/stuff/healthcheck.rb.html
=end

require 'net/http'
require 'net/https'
require 'optparse'
require 'ostruct'
require 'uri'
require 'timeout'
require 'digest/sha1'
require 'digest/md5'

Settings = OpenStruct.new
Settings.supported_hashes = %w[SHA1 MD5]

# Prepare commandline options parser
opts = OptionParser.new do |opts|
    opts.banner = "Usage: #{File.basename($0)} [options] <url>+"
    opts.separator "You must specify full URL as url"
    opts.separator "Available options:"

    opts.on('-c', '--connect-ip IP', String,
            'IP to connect to [default: from url]') do |ip|
        raise OptionParser::InvalidArgument, \
            "invalid ip" unless ip =~ /^(\d+\.){3}\d+$/
        Settings.connect_ip = ip
    end
    
    opts.on('-p', '--connect-port PORT', Integer,
            'Port to connect to [default: from url]') do |port|
        Settings.connect_port = port
    end
    
    opts.on('-s', '--source-ip IP', String,
            'IP to connect from [default: autoselect]') do |ip|
        raise OptionParser::InvalidArgument, \
            "invalid ip" unless ip =~ /^(\d+\.){3}\d+$/
        Settings.source_ip = ip
    end
    
    
    opts.on('-t', '--timeout TO', Float,
            'Timeout on connection [default: none]') do |to|
        Settings.timeout = to
    end

    opts.on('-b', '--show-body', 'Show body, not digests') do
        raise OptionParser::InvalidArgument, \
            "can't show both body and digest(s)" if Settings.digests
        raise OptionParser::InvalidArgument, \
            "can't show body and run benchmark" if Settings.benchmark
        Settings.show_body = true
    end

    opts.on('-d', '--show-digest DIGEST', String,
            "Show digest [default: #{Settings.supported_hashes.join(', ')}]",
            "May be given multiple times") do |dig|
        raise OptionParser::InvalidArgument, \
            "can't show both body and digest(s)" if Settings.show_body
        raise OptionParser::InvalidArgument, \
            "can't show digest(s) and run benchmark" if Settings.benchmark
        dig.upcase!
        raise OptionParser::InvalidArgument, \
            "invalid digest: #{dig}" unless Settings.supported_hashes.include?(dig)
        Settings.digests ||= []
        Settings.digests << dig unless Settings.digests.include?(dig)
    end

    opts.on('-v', '--verbose', "Be verbose, display progress") do
        Settings.verbose = true
    end

    opts.on('-B', '--benchmark', "Display benchmark") do
        raise OptionParser::InvalidArgument, \
            "can't show both body and digest(s)" if Settings.show_body
        raise OptionParser::InvalidArgument, \
            "can't show digest(s) and run benchmark" if Settings.benchmark
        Settings.benchmark = true
    end

    opts.on('-E', '--error-text ET', String,
            "What to display on benchmark error",
            "[default: ERROR|TIMEOUT|BODY]") do |be|
        Settings.error_text = be
    end

    opts.on('-C', '--check-body CB', String,
            "What to check body against during benchmark") do |cb|
        Settings.check_body = cb
    end

    opts.on_tail('-h', '--help', 'Show this message') do
        $stderr.puts opts
        exit 0
    end
end

# Parse commandline options
begin
    opts.parse!(ARGV)
    raise "no url given (--help for help)" if ARGV.empty?
    Settings.urls = []
    wrong_urls = []
    for url in ARGV
        uri = URI.parse(url)
        if %w[http https].include?(uri.scheme)
            Settings.urls << uri
        else
            wrong_urls << url
        end
    end
    raise "invalid URL(s): " + wrong_urls.join(' ') unless wrong_urls.empty?
rescue
    $stderr.puts "Error: #{$!}"
    exit 1
end

TestResponse = Struct.new(:status, :body, :digests, :time)

# Main test routine that's responsible for the http(s) request
def run_test(uri)
    uri = URI.parse(uri) unless uri.kind_of?(URI)
    ret = TestResponse.new
    ret.status = :ok

    # setup connection
    unless %w[http https].include?(uri.scheme)
        return [:error, "Unsupported URL scheme: #{uri.scheme}"]
    end

    http = Net::HTTP.new(Settings.connect_ip || uri.host, 
            Settings.connect_port || uri.port)
    http.local_host = Settings.source_ip if Settings.source_ip

    if uri.scheme == "https"  # enable SSL/TLS
        http.use_ssl = true
        http.verify_mode = OpenSSL::SSL::VERIFY_NONE
    end

    # Prepare worker that will probe via HTTP
    worker = proc do
        http.start do
            http.instance_variable_set("@address", uri.host)
            http.instance_variable_set("@port", uri.port)
            q = uri.path.empty? ? "/" : uri.path
            q += "?" + uri.query unless uri.query.nil?
            http.request_get(q) do |res|
                ret.body = res.body
                ret.digests = {}
                for d in Settings.supported_hashes
                        ret.digests[d] = Digest.const_get(d).hexdigest(res.body)
                end
            end
        end
    end

    # Execute worker with or without timeout, catching all errors
    start_time = nil
    end_time = nil
    begin
        start_time = Time.now
        if Settings.timeout
            timeout(Settings.timeout, &worker)
        else
            worker.call
        end
        end_time = Time.now
    rescue Timeout::Error
        ret.status = :timeout
        ret.body = "Timeout!"
    rescue
        ret.status = :error
        ret.body = "Error: #{$!} (#{$!.class})"
    rescue Interrupt
        ret.status = :interrupted
        ret.body = "Interrupted."
    end

    if start_time && end_time
        ret.time = ((end_time - start_time)*1000).to_i
    end

    ret
end

bench = []
all_ok = true
# For each url requested ...
for url in Settings.urls
    $stderr.puts "Testing url: #{url} ..." if Settings.verbose

    # Run test, remember status
    resp = run_test(url)
    all_ok &&= (resp.status == :ok)
    exit 1 if resp.status == :interrupted

    # Display results based on config
    if Settings.benchmark
        if resp.status == :ok
            if !Settings.check_body || resp.body.strip == Settings.check_body.strip
                bench << resp.time
            else
                bench << (Settings.error_text || "BODY")
                all_ok = false # body not equal, we die in pain
            end
        else
            bench << (Settings.error_text || resp.status.to_s.upcase)
        end
    elsif resp.status == :ok
        if Settings.show_body
            print resp.body
        else
            Settings.digests ||= Settings.supported_hashes
            if Settings.digests.size == 1
                puts resp.digests[Settings.digests.first]
            else
                for d in Settings.digests
                    puts "#{d} = #{resp.digests[d]}"
                end
            end
        end
    else
        puts resp.body
    end
    $stdout.flush
    if Settings.verbose
        $stderr.puts "" unless resp.body[-1] == ?\n || Settings.benchmark
        $stderr.puts "Took: #{resp.time}"
        $stderr.puts ""
    end
end

# Display benchmark at the end
unless bench.empty?
    puts bench.map { |x| x.to_s }.join(":")
end

# Exit accordingly
exit(all_ok ? 0 : 1)