#!/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.8 2001/11/15 11:34:45 akira Exp $
#
# Homepage: <URL:http://arika.org/ruby/sshpfw.html>
# Download: <URL:http://arika.org/archive/sshpfw.rb>
#

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

class Launcher
  include Observable

  class Error < StandardError; end
  class ConnectionTimeout < Error; end
  class ConnectionKilled  < Error; end
  class CommandExecFailed < Error; end

  STATUS = [:not_connect, :connecting, :connected]

  def initialize(command, interval = 30, timeout = 10, retry_times = 3)
    @command     = command
    @interval    = interval
    @timeout     = timeout
    @retry_times = retry_times

    @debug_label = @command.split(/\s+/).last if $DEBUG

    @mutex  = Mutex.new
    @status = :not_connect

    @rd = @wr = @pid = nil
    @queue  = Queue.new
    @write_th = Thread.new do
      loop do
	begin
	  dputs "write_th: waiting... " if $DEBUG
	  Thread.stop
	  sleep 3 # XXX
	  loop do
	    test_num = Time.now.to_i.to_s

	    dputs "write_th: test_num = #{test_num}" if $DEBUG
	    @wr.puts "echo #{test_num}"
	    @queue.push test_num

	    dputs "write_th: sleep(#{interval})... " if $DEBUG
	    sleep(interval)
	  end

	rescue Error
	  # the exception came from @main_th
	  dputs "write_th: killed by main_th" if $DEBUG

	rescue TimeoutError
	  # sent a command but could not get any reply.
	  queue.clear
	  dputs "write_th: queue are cleared" if $DEBUG

	rescue Exception => e
	  # socket closed?
	  dputs "write_th: socket closed?: #{$!}\n\t#{$@.join(\"\t\n\")}" if $DEBUG
	  @main_th.raise e, 'connection are closed'
	end
      end # loop

      raise
    end # @write_th

    @main_th = Thread.new do
      loop do
	begin
	  dputs "main_th: waiting... " if $DEBUG
	  Thread.stop
	  @mutex.synchronize do
	    begin
	      pre_launch
	      launch

	    rescue Error
	      # noop

	    ensure
	      if @pid
		begin
		  Process.kill(0, @pid)
		  Process.waitpid(@pid) rescue Errno::ECHILD
		rescue Errno::ESRCH
		end
	      end
	      @rd = @wr = @pid = nil
	      post_launch
	    end
	  end

	rescue RuntimeError
	  # command exec failed
	  dputs "main_th: exec failed" if $DEBUG
	  @write_th.raise Error
	  raise CommandExecFailed, 'exec failed'

	rescue Exception => e
	  # socket closed?
	  dputs "main_th: socket closed?: #{$!}\n#{$@.join(\"\n\")}" if $DEBUG
	  raise
	end
      end # loop

      raise
    end # @main_th
  end # initialize
  attr_writer :command, :interval, :timeout, :retry_times
  attr_reader :status

  private

  def dputs(msg)
    $stderr.puts 'D: ' + 
      Time.now.strftime('%H:%M:%S: ') + 
      @debug_label + ': ' + msg.to_s
  end

  def update_status(status)
    if STATUS.include?(status)
      @status = status
    else
      raise ArgumentError, "unknown status: #{status}"
    end

    changed
    notify_observers(self, @status)
  end

  def pre_launch
  end

  def pre_connect
    update_status(:connecting)
  end

  def post_connect
    update_status(:connected)
  end

  def write_and_end
    @wr.puts 'exit'
    dputs "write_and_end: " if $DEBUG
  end

  def last_read
    tmp = @rd.sysread(1024)
    dputs "last_read: #{tmp}" if $DEBUG && tmp
  end

  def post_launch
    update_status(:not_connect)
  end

  def launch
    retry_times = @retry_times

    PTY.spawn(@command) do |@rd, @wr, @pid|
      pre_connect

      @write_th.run
      begin
	test = nil
	loop do
	  str = nil
	  begin
	    timeout(@timeout) {str = @rd.gets}
	    dputs "launch: str = #{str}" if $DEBUG

	    if !test
	      # get a command which ware sent by write_th.
	      test = @queue.pop(true) rescue ThreadError

	    else
	      if str.index(test) == 0
		dputs "launch: #{test} (#{Time.at(test.to_i).strftime('%H:%M:%S')}) OK" if $DEBUG
		if @status != :connected
		  post_connect
		  @status = :connected
		end
		test = nil
	      end
	    end

	  rescue TimeoutError
	    if test
	      # command ware sent but no reply was gotten.
	      # -> connection timed out
	      dputs "launch: Timed out (wait: #{test})" if $DEBUG
	      retry_times -= 1
	      if retry_times > 0
		dputs "launch: retrying..." if $DEBUG
		@write_th.raise TimeoutError
		test = str = nil
	      else
		raise
	      end

	    else
	      retry
	    end
	  end

	end # loop

      rescue TimeoutError
	# timed out
	dputs "launch: connection aborted" if $DEBUG
	raise ConnectionTimeout, "no response from the login host"

      rescue SystemCallError, IOError
	# socket closed?
	dputs "launch: connection are closed" if $DEBUG
	raise ConnectionKilled, 'connection are closed'

      rescue RuntimeError
	# connectin closed
	dputs "launch: connection are killed" if $DEBUG
	raise ConnectionKilled, "connection are killed"

      rescue Error
	# user disconnect
	dputs "launch: user disconnect" if $DEBUG

      rescue => e
	# unknown error
	dputs "launch: unknown error occurred" if $DEBUG
	raise e, "unknown error occurred"

      ensure
	write_and_end rescue RuntimeError
	last_read     rescue RuntimeError
	
	@rd.close
	@wr.close
      end
    end # PTY.spawn

    dputs "launch: finished" if $DEBUG
  end # launch

  public

  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 connected?
    status_is?(:connected)
  end

  def connect
    @main_th.run
  end

  def disconnect(wait = false)
    @main_th.raise Error, 'user disconnect'
    wait_disconnect if wait
  end

  def wait_disconnect
    # wait for closing connection
    @mutex.synchronize {}
  end

  def inspect
    sprintf('#<%s:0x%x>', self.class, self.id)
  end
end # Launcher

class SSHConnection < Launcher
  def initialize(host, args = [])
    @host = host
    @args = args

    super(['ssh', args, host].flatten.join(' '))
  end
  attr_reader :host, :args
end # SSHConnection

class SSHConnectionButton < Gtk::ToggleButton
  STATUS = {
    :not_connect => '>   <',
    :connecting  => '>- -<',
    :connected   => '>---<',
  }

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

    @host = host
    @args = args

    @ssh  = SSHConnection.new(host, args)
    @ssh.add_observer(self)

    @label = Gtk::Label.new(STATUS[@ssh.status])
    self.add(@label)

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

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

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

  def close_connection
    @ssh.disconnect(true)
  end

  # for Observable
  def update(ssh, status)
p [@host, ssh, status]
    @label.set(STATUS[status])
  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

if $0 == __FILE__
  #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"
  rescue RuntimeError
    # PTY raises RuntimeError on main(or current?) thread if command failed
    p [666, $!, $@]
  ensure
p [666]
    buttons.each do |button|
      button.close_connection
    end
    p [:end, $!, $@]
  end

  exit 0
end