#!/usr/bin/env ruby
################################################################
#
# rcontool - shell interface to rcon commands
#
# (C) 2006 Erik Hollensbe, License details below
#
# Use 'rcontool -h' for usage instructions.
#
# The compilation of software known as rcontool is distributed under the
# following terms:
# Copyright (C) 2005-2006 Erik Hollensbe. All rights reserved.
#
# Redistribution and use in source form, with or without
# modification, are permitted provided that the following conditions
# are met:
# 1. Redistributions of source code must retain the above copyright notice, this list of conditions and the following disclaimer.
#
# THIS SOFTWARE IS PROVIDED BY AUTHOR AND CONTRIBUTORS ``AS IS'' AND
# ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
# ARE DISCLAIMED.  IN NO EVENT SHALL AUTHOR OR CONTRIBUTORS BE LIABLE
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS
# OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION)
# HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
# LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY
# OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF
# SUCH DAMAGE.
#
#
################################################################

#
# rubygems hack
#

begin
  require 'rubygems'
rescue LoadError => e
end
begin
  require 'rcon'
  require 'ip'
rescue LoadError => e
  $stderr.puts "rcontool requires the rcon and ip libraries be installed."
  $stderr.puts "You can find them both via rubygems or at http://rubyforge.org."
  exit -1
end

RCONTOOL_VERSION = '0.1.0'

require 'optparse'
require 'ostruct'

#
# Manages our options
#

def get_options
  options = OpenStruct.new
  # ip address (IP::Address object)
  options.ip_address = nil
  # port (integer)
  options.port = nil
  # password
  options.password = nil
  # protocol type (one of :hlds, :source, :oldquake, :newquake)
  options.protocol_type = nil
  # verbose, spit out extra information
  options.verbose = false
  # command to execute on the server
  options.command = nil

  optparse = OptionParser.new do |opts|
    opts.banner = "Usage: #{File.basename $0} <ip_address:port> <command> [options]"
    opts.separator ""
    opts.separator "Options:"

    opts.on("--ip-address [ADDRESS]",
            "Provide an IP address to connect to. Does not take a port.") do |ip_address|
      if ! options.ip_address.nil?
        $stderr.puts "Error: you have already provided an IP Address."
        $stderr.puts opts
        exit -1
      end

      options.ip_address = IP::Address.new(ip_address)
    end

    opts.on("-r", "--port [PORT]",
            "Port to connect to.") do |port|
      if ! options.port.nil?
        $stderr.puts "Error: you have already provided a port."
        $stderr.puts opts
        exit -1
      end

      options.port = port.to_i
    end

    opts.on("-c", "--command [COMMAND]",
            "Command to run on the server.") do |command|
      if ! options.command.nil?
        $stderr.puts "Error: you have already provided a command."
        $stderr.puts opts
        exit -1
      end
      
      options.command = command
    end

    opts.on("-p", "--password [PASSWORD]",
            "Provide a password on the command line.") do |password|
      options.password = password
    end

    opts.on("-f", "--password-from [FILENAME]",
            "Get the password from a file (use '/dev/fd/0' or '/dev/stdin' to read from Standard Input).") do |filename|
      if !filename.nil?
        f = File.open(filename)
        options.password = f.gets.chomp
        f.close
      else
        $stderr.puts "Error: filename (from -f) is not valid."
        $stderr.puts opts
        exit -1
      end
    end

    opts.on("-t", "--protocol-type [TYPE]", [:hlds, :source, :oldquake, :newquake],
            "Type of rcon connection to make: (hlds, source, oldquake, newquake).",
            "   Note: oldquake is quake1/quakeworld, newquake is quake2/3.") do |protocol_type|
      options.protocol_type = protocol_type
    end

    opts.on("-v", "--[no-]verbose", 
            "Run verbosely, print information about each packet recieved and turnaround times.") do |verbose|
      options.verbose = verbose
    end

    opts.on("-h", "--help",
                 "This help message.") do
      $stderr.puts opts
      exit -1
    end

    opts.on("--version", "Print the version information.") do
      $stderr.puts "This is rcontool version #{RCONTOOL_VERSION},"
      $stderr.puts "it is located at #{File.expand_path $0}."
      exit -1
    end

    opts.separator ""
    opts.separator "Note: IP, port, protocol type, password and command are required to function."
    opts.separator ""
    opts.separator "Examples (all are equivalent):"
    opts.separator "\t#{File.basename($0)} 10.0.0.11 status -t hlds -r 27015 -p foobar"
    opts.separator "\techo 'foobar' | #{File.basename($0)} 10.0.0.11:27015 status -t hlds -f /dev/stdin"
    opts.separator "\t#{File.basename($0)} --ip-address 10.0.0.11 --port 27015 -c status -t hlds -f file_with_password"
    opts.separator ""

  end

  ################################################################
  #
  # This hackery is to help facilitate the bareword options if
  # they exist, while still allowing for the option parser
  # to work properly.
  #
  ################################################################

  s1 = ARGV.shift
  s2 = ARGV.shift

  begin
    options.ip_address = IP::Address::IPv4.new(s1)
    options.command = s2
  rescue IP::AddressException => e
    # attempt to split it first... not sure how to best handle this situation
    begin
      ip,port = s1.split(/:/, 2)
      options.ip_address = IP::Address::IPv4.new(ip)
      options.port = port.to_i
      options.command = s2
    rescue Exception => e
    end

    if [options.ip_address, options.port].include? nil
      ARGV.unshift(s2)
      ARGV.unshift(s1)
    end
  end
  
  optparse.parse!

  if [options.ip_address, options.protocol_type, options.port, options.password, options.command].include? nil
    $stderr.puts optparse
    exit -1
  end

  return options
end

def verbose(string)
  $stderr.puts string if $options.verbose
end

def dump_source_packet(packet)
  if $options.verbose
    verbose "Request ID: #{packet.request_id}"
    verbose "Packet Size: #{packet.packet_size}"
    verbose "Response Type: #{packet.command_type}"
  end
end

################################################################
#
# start main block
#
################################################################

$options = get_options

################################################################
#
# Source query
#
################################################################

if $options.protocol_type == :source
  verbose "Protocol type 'SOURCE' selected."

  rcon = RCon::Query::Source.new($options.ip_address.ip_address, $options.port)

  # if we have a verbose request, give all the information we can about
  # the query, including the packet information.
  rcon.return_packets = $options.verbose

  verbose "Attempting authentication to #{$options.ip_address.ip_address}:#{$options.port} with password '#{$options.password}'"

  value = rcon.auth $options.password

  dump_source_packet value

  if ($options.verbose && value.command_type == RCon::Packet::Source::RESPONSE_AUTH) || value
    verbose "Authentication succeeded. Sending command: '#{$options.command}'"

    value = rcon.command $options.command

    dump_source_packet value
    verbose ""

    if $options.verbose
      puts value.string1
    else
      puts value
    end
    
    exit 0
  else
    $stderr.puts "Authentication failed."
    exit 1
  end

################################################################
#
# Original Query
#
################################################################

else
  rcon = nil
  case $options.protocol_type
  when :hlds
    verbose "Protocol type 'HLDS' selected"
    rcon = RCon::Query::Original.new($options.ip_address.ip_address, $options.port, $options.password,
                                     RCon::Query::Original::HLDS)
  when :oldquake
    verbose "Protocol type 'OLDQUAKE' selected"
    rcon = RCon::Query::Original.new($options.ip_address.ip_address, $options.port, $options.password,
                                     RCon::Query::Original::QUAKEWORLD)
  when :newquake
    verbose "Protocol type 'NEWQUAKE' selected"
    rcon = RCon::Query::Original.new($options.ip_address.ip_address, $options.port, $options.password,
                                     RCon::Query::Original::NEWQUAKE)
  end
  verbose "Attempting transmission to #{$options.ip_address.ip_address}:#{$options.port}"
  verbose "Using password: '#{$options.password}' and sending command: '#{$options.command}'"
  verbose ""
  string = rcon.command($options.command)

  puts string
  exit 0
end