lib/net/ssh/shell.rb in net-ssh-shell-0.1.0 vs lib/net/ssh/shell.rb in net-ssh-shell-0.2.0

@@ -9,25 +9,29 @@
       attr_reader :session
       attr_reader :channel
       attr_reader :state
       attr_reader :shell
       attr_reader :processes
+      attr_accessor :default_process_class
 
       def initialize(session, shell=:default)
         @session = session
         @shell = shell
         @state = :closed
         @processes = []
         @when_open = []
+        @on_process_run = nil
+        @default_process_class = Net::SSH::Shell::Process
         open
       end
 
       def open(&callback)
         if closed?
           @state = :opening
           @channel = session.open_channel(&method(:open_succeeded))
           @channel.on_open_failed(&method(:open_failed))
+          @channel.on_request('exit-status', &method(:on_exit_status))
         end
         when_open(&callback) if callback
         self
       end
 
@@ -58,25 +62,37 @@
 
       def opening?
         !open? && !closed?
       end
 
-      def execute(command, klass=Net::SSH::Shell::Process, &callback)
-        process = klass.new(self, command, callback)
-        process.run if processes.empty?
+      def on_process_run(&callback)
+        @on_process_run = callback
+      end
+
+      def execute(command, *args, &callback)
+        # The class is an optional second argument.
+        klass = default_process_class
+        klass = args.shift if args.first.is_a?(Class)
+
+        # The properties are expected to be the next argument.
+        props = {}
+        props = args.shift if args.first.is_a?(Hash)
+
+        process = klass.new(self, command, props, callback)
         processes << process
-        return process
+        run_next_process if processes.length == 1
+        process
       end
 
       def subshell(command, &callback)
         execute(command, Net::SSH::Shell::Subshell, &callback)
       end
 
       def execute!(command, &callback)
-        execute(command, &callback)
+        process = execute(command, &callback)
         wait!
-	return process
+        process
       end
 
       def busy?
         opening? || processes.any?
       end
@@ -90,11 +106,11 @@
       end
 
       def child_finished(child)
         channel.on_close(&method(:on_channel_close)) if !channel.nil?
         processes.delete(child)
-        processes.first.run if processes.any?
+        run_next_process
       end
 
       def separator
         @separator ||= begin
           s = Digest::SHA1.hexdigest([session.object_id, object_id, Time.now.to_i, Time.now.usec, rand(0xFFFFFFFF)].join(":"))
@@ -107,45 +123,59 @@
         @channel = nil
       end
 
       private
 
-        def open_succeeded(channel)
-          @state = :pty
-          channel.on_close(&method(:on_channel_close))
-          channel.request_pty(:modes => { Net::SSH::Connection::Term::ECHO => 0 }, &method(:pty_requested))
+      def run_next_process
+        if processes.any?
+          process = processes.first
+          @on_process_run.call(self, process) if @on_process_run
+          process.run
         end
+      end
 
-        def open_failed(channel, code, description)
-          @state = :closed
-          raise "could not open channel for process manager (#{description}, ##{code})"
-        end
+      def open_succeeded(channel)
+        @state = :pty
+        channel.on_close(&method(:on_channel_close))
+        channel.request_pty(:modes => { Net::SSH::Connection::Term::ECHO => 0 }, &method(:pty_requested))
+      end
 
-        def pty_requested(channel, success)
-          @state = :shell
-          raise "could not request pty for process manager" unless success
-          if shell == :default
-            channel.send_channel_request("shell", &method(:shell_requested))
-          else
-            channel.exec(shell, &method(:shell_requested))
-          end
+      def open_failed(channel, code, description)
+        @state = :closed
+        raise "could not open channel for process manager (#{description}, ##{code})"
+      end
+
+      def on_exit_status(channel, data)
+        unless data.read_long == 0
+          raise "the shell exited unexpectedly"
         end
+      end
 
-        def shell_requested(channel, success)
-          @state = :initializing
-          raise "could not request shell for process manager" unless success
-          channel.on_data(&method(:look_for_initialization_done))
-          channel.send_data "export PS1=; echo #{separator} $?\n"
+      def pty_requested(channel, success)
+        @state = :shell
+        raise "could not request pty for process manager" unless success
+        if shell == :default
+          channel.send_channel_request("shell", &method(:shell_requested))
+        else
+          channel.exec(shell, &method(:shell_requested))
         end
+      end
 
-        def look_for_initialization_done(channel, data)
-          if data.include?(separator)
-            @state = :open
-            @when_open.each { |callback| callback.call(self) }
-            @when_open.clear
-          end
+      def shell_requested(channel, success)
+        @state = :initializing
+        raise "could not request shell for process manager" unless success
+        channel.on_data(&method(:look_for_initialization_done))
+        channel.send_data "export PS1=; echo #{separator} $?\n"
+      end
+
+      def look_for_initialization_done(channel, data)
+        if data.include?(separator)
+          @state = :open
+          @when_open.each { |callback| callback.call(self) }
+          @when_open.clear
         end
+      end
     end
   end
 end
 
 class Net::SSH::Connection::Session
@@ -153,8 +183,8 @@
   # session. Yields the new shell if a block is given. Returns the shell
   # instance.
   def shell(*args)
     shell = Net::SSH::Shell.new(self, *args)
     yield shell if block_given?
-    return shell
+    shell
   end
 end