#!/usr/bin/ruby
#
# Copyright (c) 2001 akira yamada <akira@ruby-lang.org>
# You can redistribute it and/or modify it under the same term as Ruby.
#  
# $Id: sshpfw.rb,v 1.4 2001/11/11 17:37:03 akira Exp $
#
# ToDo:
#  - show error message
#  - re-load configuration file
#  - auto re-connect
#  - panel-applet
#  - GUI for configuration
#

require 'pty'
require 'thread'
require 'timeout'
require 'gtk'

module Launcher
  class ConnectionTimeout < StandardError; end
  class ConnectionKilled  < StandardError; end
  class CommandExecFailed < StandardError; end

  def pre_launch
  end

  def pre_connect
  end

  def post_connect
  end

  def post_launch
  end

  def launch(cmd, interval = 30, timeout = 10, retry_times = 3)
    queue = Queue.new
    cmd   = if cmd.kind_of?(Array)
	      cmd.join(' ')
	    else
	      cmd.to_s
	    end

    connected = false

    pre_launch if respond_to?(:pre_launch)
    begin
      PTY.spawn(cmd) do |rd, wr, pid|
	begin
	  pre_connect if respond_to?(:pre_connect)

	  th = Thread.new do
	    begin
	      loop do
		test_num = Time.now.to_i.to_s
		wr.puts "echo #{test_num}"
		queue.push test_num

		sleep interval
	      end

	    rescue TimeoutError
	      # sent a command but could not get any reply.
	      queue.clear
	      retry
	    end
	  end

	  test = nil
	  loop do
	    str = nil
	    begin
	      timeout(timeout) {str = rd.gets}
	    rescue TimeoutError
	      if test
		# command ware sent but no reply was gotten.
		# -> connection timed out
		retry_times -= 1
		if retry_times > 0
		  th.raise
		  retry
		else
		  raise
		end

	      else
		retry
	      end
	    end

	    if test && str
	      if str.index(test)
		$stderr.puts "D: #{cmd}: #{Time.at(test.to_i)} OK" if $DEBUG
		unless connected
		  post_connect  if respond_to?(:post_connect)
		  connected = true
		end
		test = nil
	      end

	    else
	      begin
		# get a command which ware sent by th.
		test = queue.pop(true)
	      rescue ThreadError
		# queue is empty
	      end
	    end
	  end # loop

	rescue TimeoutError
	  # timed out
	  $stderr.puts "D: #{cmd}: timed out" if $DEBUG
	  raise ConnectionTimeout, "timed out"

	rescue RuntimeError
	  # connectin closed
	  $stderr.puts "D: #{cmd}: connection are killed" if $DEBUG
	  raise ConnectionKilled, "connection are killed"

	rescue => e
	  # unknown error
	  $stderr.puts "D: #{cmd}: unknown error occurred" if $DEBUG
	  raise e, "unknown error occurred"

	ensure
	  begin
	    th.exit
	    rd.close
	    wr.close
	  rescue Exception
	  end
	end
      end # PTY.spawn

    rescue RuntimeError
      # cmd exec failed
      $stderr.puts "D: #{cmd}: exec failed" if $DEBUG
      raise CommandExecFailed, "exec failed"

    ensure
      post_launch if respond_to?(:post_launch)
    end

    $stderr.puts "D: #{cmd}: finished" if $DEBUG
  end
  module_function :launch
end # Launcher

class SSHConnectionButton < Gtk::ToggleButton
  include Launcher

  Status = {
    :not_connect => ['>   <', false],
    :connecting  => ['>- -<', true],
    :connected   => ['>---<', true],
  }

  def initialize(host, args = [])
    super()

    @host = host
    @args = args
    @thread = nil
    @status = :not_connect
    @label = Gtk::Label::new(Status[@status].first)
    self.add(@label)

    self.signal_connect('button_press_event') {|btn, ev|
    }
    self.signal_connect('button_release_event') {|btn, ev|
      if ev.button == 1
	connect
      elsif ev.button == 3
	# config
      end
    }
    self.signal_connect('clicked') {|arg|
      self.active = Status[@status].last
    }
  end

  def set_label(str)
    @label.set(str)
  end

  def get_label
    @label.get
  end

  def show(*arg)
    @label.show
    super(*arg)
  end

  def hide(*arg)
    super(*arg)
  end

  def status_is?(status)
    if Status.include?(status)
      if @status == status
	return true
      else
	return false
      end
    else
      raise ArgumentError, "unknown status: #{status}"
    end
  end

  def update_status(status)
    if Status.include?(status)
      @status = status
      set_label(Status[@status].first)
      unless self.active == Status[@status].last
	self.active = Status[@status].last
      end
    else
      raise ArgumentError, "unknown status: #{status}"
    end
  end

  def thread_exit
    if @thread && @thread.alive?
      @thread.exit
      update_status(:not_connect)
    end
  end

  # for Launcher

  def post_connect
    update_status(:connected)
  end

  def connect
    if status_is?(:not_connect)
      update_status(:connecting)
      @thread = Thread.new do
	begin
	  launch(['ssh', @args, @host].flatten)

	rescue ConnectionTimeout, ConnectionKilled
	  $stderr.puts "#{@host}: #{$!}"
	rescue
	  $stderr.puts "#{@host}: #{$!}"
	ensure
	  update_status(:not_connect)
	end
      end

    elsif status_is?(:connected)
      thread_exit
    else

    end
  end
end # SSHConnectionButton

#
# format 1:
#   local-port  [remote-host]:remote-port  [login-host  [options]]
# uses last value of login-host if login-host are omitted.
#
# format 2:
#   -  -  login-host  [options]
# not specify for configuration about port forwarding.
# you must configure on .ssh/config.
#
# exmaple:
#   10025 :25  foo.example.com
#   10110 :110
#   20143 :143 foo.example.jp
#   20025 :25
#   -     -    bar.example.jp
# it realizes 3 buttons:
# 1) ssh -L10025:foo.example.com:25 -L10110:foo.example.com:110 foo.example.com
# 2) ssh -L20143:foo.example.jp:143 -L20025:foo.example.jp:25 foo.example.jp
# 3) ssh bar.example.jp
#
def read_config(path)
  path = File.expand_path(path)
  return nil unless FileTest.exist?(path)

  config = {}
  lprt = rhst = rprt = svr = nil
  open(path, 'r') do |io|
    io.each do |l|
      l.chomp!
      l.sub!(/#.*$/o, '')
      l.strip!

      #    lport   rhost          rport      user    server      options
      if /^(\d+)\s+([-.0-9a-z]+)?:(\d+)(?:\s+(\w+@)?([-.0-9a-z]+)(?:\s+(.+))?)?$/io =~ l
	lprt = $1
	rprt = $3
	svr  = $5 if $5
	raise 'login-host is not specified' unless svr

	rhst = $2 ? $2 : svr
	svr  = $4 + svr if ($4 && svr)

	config[svr] ||= []
	config[svr] << $6 if $6
	config[svr] << '-L' + lprt + ':' + rhst + ':' + rprt

	#             server             options
      elsif /^-\s+-\s+((?:\w+@)?[-.0-9a-z]+)(?:\s+(.+))?$/ =~ l
	config[$1] ||= []
	config[svr] << $2 if $2
      end
    end
  end

  config.collect do |svr, args|
    [svr, args.join(' ')]
  end
end

#Thread.abort_on_exception = true

# load configuration files...
conffile = '~/.sshpfwrc'
hosts = read_config(conffile)
raise "no config file (#{conffile})" unless hosts
raise "no valide line in #{conffile}" if hosts.empty?

# load gtk resource file...
rc_file = '.gtkrc'
rc_search_path = ['.', '~']

rc_search_path.each do |path|
  f = File.join(File.expand_path(path), rc_file)
  if FileTest.file?(f)
    Gtk::RC.parse(f)
    break
  end
end

# main...
top  = Gtk::Window.new
tips = Gtk::Tooltips.new
hbox = Gtk::HBox.new

top.signal_connect('delete_event') {|*x| false}
top.signal_connect('destroy') {|*x| exit}
top.add(hbox)
hbox.show

buttons = []
hosts.each do |host, args|
  btn = SSHConnectionButton.new(host, args)
  buttons << btn

  tips.set_tip(btn, host, nil)
  hbox.pack_start(btn)
  btn.show
end

top.show
           
begin
  Gtk.main
rescue Interrupt
  $stderr.puts "interrupted"
ensure
  buttons.each do |button|
    button.thread_exit
  end
end

exit 0