lib/infrataster/server.rb in infrataster-0.1.4 vs lib/infrataster/server.rb in infrataster-0.1.5

@@ -1,9 +1,10 @@
 require 'tmpdir'
 require 'net/ssh'
 require 'net/ssh/gateway'
 require 'ipaddr'
+require 'shellwords'
 
 module Infrataster
   class Server
     Error = Class.new(StandardError)
 
@@ -82,60 +83,68 @@
         [port, finalize_proc]
       end
     end
 
     def open_gateway_on_from_server(port)
-      if from
-        new_port, finalize_proc = from.gateway_open(@address, port)
-        Logger.debug("tunnel: localhost:#{new_port} -> #{from.address} -> #{@address}:#{port}")
-        ['127.0.0.1', new_port, finalize_proc]
+      if block_given?
+        if from
+          from.gateway_open(@address, port) do |new_port|
+            Logger.debug("tunnel: localhost:#{new_port} -> #{from.address} -> #{@address}:#{port}")
+            yield '127.0.0.1', new_port
+          end
+        else
+          yield @address, port
+        end
       else
-        [@address, port, nil]
-      end
-    end
-
-    def gateway_on_from_server(port)
-      if from
-        from.gateway_open(@address, port) do |new_port|
+        if from
+          new_port, finalize_proc = from.gateway_open(@address, port)
           Logger.debug("tunnel: localhost:#{new_port} -> #{from.address} -> #{@address}:#{port}")
-          yield '127.0.0.1', new_port
+          ['127.0.0.1', new_port, finalize_proc]
+        else
+          [@address, port, nil]
         end
-      else
-        yield @address, port
       end
     end
 
     def ssh_start_args
       @ssh_start_args ||= _ssh_start_args
     end
 
     private
 
-
     def _ssh_start_args
       config = {}
+
       if @options[:ssh]
         config = @options[:ssh]
+        config[:host] ||= @address
       elsif @options[:vagrant]
-        if @options[:vagrant] != true
-          vagrant_name = @options[:vagrant]
-        else
-          vagrant_name = @name
-        end
-
-        Dir.mktmpdir do |dir|
-          output = File.join(dir, 'ssh-config')
-          `/usr/bin/vagrant ssh-config #{vagrant_name} > #{output}`
-          if $?.exitstatus != 0
-            raise Error, "`vagrant ssh-config` failed. Please check if VMs are running or not."
-          end
-          config = Net::SSH::Config.for(@name.to_s, [output])
-        end
+        vagrant_name = if @options[:vagrant] != true
+                         @options[:vagrant]
+                       else
+                         @name
+                       end
+        config = ssh_config_for_vagrant(vagrant_name)
       else
         raise Error, "Can't get configuration to connect to the server via SSH. Please set ssh option or vagrant option."
       end
 
       [config[:host], config[:user], config]
+    end
+
+    def ssh_config_for_vagrant(name)
+      config = nil
+
+      Dir.mktmpdir do |dir|
+        output = File.join(dir, 'ssh-config')
+        `/usr/bin/vagrant ssh-config #{Shellwords.shellescape(name)} > #{Shellwords.shellescape(output)}`
+        if $?.exitstatus != 0
+          raise Error, "`vagrant ssh-config` failed. Please check if VMs are running or not."
+        end
+        config = Net::SSH::Config.for(@name.to_s, [output])
+      end
+
+      config
     end
   end
 end