lib/remote/session.rb in remote-session-0.0.2 vs lib/remote/session.rb in remote-session-0.0.3

@@ -1,8 +1,10 @@
+require 'remote/session/send_file'
 require 'remote/session/version'
 require 'net/sftp'
 require 'net/ssh'
+require 'base64'
 
 module Remote
   class Session
     SUDO_PROMPT = 'sudo_prompt'
 
@@ -39,45 +41,33 @@
       raise "Session is closed" if @session.nil?
       puts "@#{ @host }: #{ command }"
       puts @session.exec!( command )
     end
 
-    def sudo( command )
+    def sudo( commands )
       raise "Session is closed" if @session.nil?
+      commands = [ *commands ]
 
-      puts "@#{ @host }: sudo #{ command }"
       @session.open_channel do |ch|
         ch.request_pty do |ch, success|
           raise "Could not obtain pty" if ! success
 
-          channel_exec ch, command
+          channel_exec ch, commands
         end
       end
       @session.loop
     end
 
     def close
       @session.close
       @session = nil
     end
 
-    def sudo_put( remote_path, &block )
-      temp_path = "/tmp/remote-session.#{ Time.now.to_f }"
-      run "mkdir #{ temp_path }"
-      run "chmod 0700 #{ temp_path }"
-
-      temp_file = File.join( temp_path, File.basename( remote_path ) )
-      put temp_file, &block
-
-      sudo "cp -f #{ temp_file } #{ remote_path }"
-      run "rm -rf #{ temp_path }"
-    end
-    
     def put( remote_path, &block )
       sftp = Net::SFTP::Session.new( @session ).connect!
       sftp.file.open( remote_path, 'w' ) do |f|
-        f.puts block.call
+        f.write block.call
       end
       sftp.close_channel
     end
 
     private
@@ -92,31 +82,72 @@
 
     def connect
       @session = Net::SSH.start( @host, @username, ssh_options )
     end
 
-    def channel_exec( ch, command )
-      ch.exec "sudo -p '#{ SUDO_PROMPT }' #{ command }" do |ch, success|
-        raise "Could not execute sudo command: #{ command }" if ! success
+    def channel_exec( ch, commands )
+      ch[ :commands ] = commands
 
-        ch.on_data do | ch, data |
-          if data =~ Regexp.new( SUDO_PROMPT )
-            ch.send_data "#{ @sudo_password }\n"
-          else
-            prompt_matched = false
-            @prompts.each_pair do | prompt, send |
-              if data =~ Regexp.new( prompt )
-                ch.send_data "#{ send }\n"
-                prompt_matched = true
-              end
-            end
-            puts data if ! prompt_matched
-          end
-        end
+      ch.exec "sudo -k -p '#{ SUDO_PROMPT }' su -" do |ch, success|
+        raise "Could not execute sudo su command" if ! success
 
+        ch.on_data( &method( :on_data ) )
+
         ch.on_extended_data do |ch, type, data|
-          raise "Error #{ data } while performing command: #{ command }"
+          $stderr.puts data
         end
+      end
+    end
+
+    def on_data( ch, data )
+      if data =~ Regexp.new( SUDO_PROMPT )
+        ch.send_data "#{ @sudo_password }\n"
+        return
+      end
+
+      @prompts.each_pair do | prompt, send |
+        if data =~ Regexp.new( prompt )
+          ch.send_data "#{ send }\n"
+          return
+        end
+      end
+
+      $stdout.write( data )
+
+      if ch[ :commands ].size == 0
+        ch.send_data "exit\n"
+        return
+      end
+
+      command = ch[ :commands ].shift
+      if command.is_a?( Remote::Session::SendFile )
+        send_file_chunk( ch, command )
+      else
+        ch.send_data "#{command}\n"
+      end
+    end
+
+    def send_file_chunk( ch, send_file )
+      if send_file.open?
+        operator = '>>'
+      else
+        send_file.open
+        operator = '>'
+      end
+
+      chunk =
+        if ! send_file.eof?
+          Base64.encode64( send_file.read )
+        else
+          # Handle empty files
+          ''
+        end
+      ch.send_data "echo -n '#{ chunk }' | base64 -d #{ operator } #{ send_file.remote_path }\n"
+
+      if send_file.eof?
+        send_file.close
+      else
+        ch[ :commands ].unshift send_file
       end
     end
 
   end
 end