From e34f0a8af7b5a0f9dcef5b2041fd5226869a4193 Mon Sep 17 00:00:00 2001
From: Mitchell Hashimoto <mitchell.hashimoto@gmail.com>
Date: Sat, 17 Dec 2011 17:16:11 -0800
Subject: [PATCH] `vagrant ssh`

---
 lib/vagrant.rb             |  1 +
 lib/vagrant/command.rb     |  1 +
 lib/vagrant/command/ssh.rb | 70 ++++++++++++++++++++++----------------
 3 files changed, 43 insertions(+), 29 deletions(-)

diff --git a/lib/vagrant.rb b/lib/vagrant.rb
index 5858a7a56..d018c3363 100644
--- a/lib/vagrant.rb
+++ b/lib/vagrant.rb
@@ -100,6 +100,7 @@ Vagrant.commands.register(:package) { Vagrant::Command::Package }
 Vagrant.commands.register(:provision) { Vagrant::Command::Provision }
 Vagrant.commands.register(:reload) { Vagrant::Command::Reload }
 Vagrant.commands.register(:resume) { Vagrant::Command::Resume }
+Vagrant.commands.register(:ssh) { Vagrant::Command::SSH }
 Vagrant.commands.register(:suspend) { Vagrant::Command::Suspend }
 Vagrant.commands.register(:up) { Vagrant::Command::Up }
 
diff --git a/lib/vagrant/command.rb b/lib/vagrant/command.rb
index 48c3cfba0..0544b51e5 100644
--- a/lib/vagrant/command.rb
+++ b/lib/vagrant/command.rb
@@ -8,6 +8,7 @@ module Vagrant
     autoload :Provision, 'vagrant/command/provision'
     autoload :Reload, 'vagrant/command/reload'
     autoload :Resume, 'vagrant/command/resume'
+    autoload :SSH, 'vagrant/command/ssh'
     autoload :Suspend, 'vagrant/command/suspend'
     autoload :Up, 'vagrant/command/up'
   end
diff --git a/lib/vagrant/command/ssh.rb b/lib/vagrant/command/ssh.rb
index bea64ebb5..03f8b811a 100644
--- a/lib/vagrant/command/ssh.rb
+++ b/lib/vagrant/command/ssh.rb
@@ -1,49 +1,61 @@
+require 'optparse'
+
 module Vagrant
   module Command
-    class SSHCommand < NamedBase
-      class_option :command, :type => :string, :default => false, :aliases => "-c"
-      register "ssh", "SSH into the currently running Vagrant environment."
-
+    class SSH < Base
       def execute
-        if options[:command]
-          ssh_execute
-        else
-          ssh_connect
-        end
-      end
+        options = {}
 
-      protected
+        opts = OptionParser.new do |opts|
+          opts.banner = "Usage: vagrant ssh [vm-name] [-c command]"
 
-      def ssh_execute
-        ssh_vm.ssh.execute do |ssh|
-          ssh.exec!(options[:command]) do |channel, type, data|
-            if type != :exit_status
-              # Print the SSH output as it comes in, but don't prefix it and don't
-              # force a new line so that the output is properly preserved
-              ssh_vm.env.ui.info(data.to_s, :prefix => false, :new_line => false)
-            end
+          opts.separator ""
+
+          opts.on("-c", "--command COMMAND", "Execute an SSH command directly.") do |c|
+            options[:command] = c
           end
         end
-      end
 
-      def ssh_connect
-        ssh_vm.ssh.connect
-      end
+        argv = parse_options(opts)
+        return if !argv
 
-      def ssh_vm
-        @ssh_vm ||= begin
-          vm = self.name.nil? && env.multivm? ? env.primary_vm : nil
-          raise Errors::MultiVMTargetRequired, :command => "ssh" if !vm && target_vms.length > 1
-          vm = target_vms.first if !vm
+        # SSH always requires a target VM
+        raise Errors::MultiVMTargetRequired, :command => "ssh" if @env.multivm? && !argv[0]
 
+        # Execute the actual SSH
+        with_target_vms(argv[0]) do |vm|
           # Basic checks that are required for proper SSH
           raise Errors::VMNotCreatedError if !vm.created?
           raise Errors::VMInaccessible if !vm.vm.accessible?
           raise Errors::VMNotRunningError if !vm.vm.running?
 
-          vm
+          if options[:command]
+            ssh_execute(vm, options[:command])
+          else
+            ssh_connect(vm)
+          end
         end
       end
+
+      protected
+
+      def ssh_execute(vm, command=nil)
+        @logger.debug("Executing command: #{command}")
+        vm.ssh.execute do |ssh|
+          ssh.exec!(command) do |channel, type, data|
+            if type != :exit_status
+              # Print the SSH output as it comes in, but don't prefix it and don't
+              # force a new line so that the output is properly preserved
+              vm.ui.info(data.to_s, :prefix => false, :new_line => false)
+            end
+          end
+        end
+      end
+
+      def ssh_connect(vm)
+        @logger.debug("`exec` into ssh prompt")
+        vm.ssh.connect
+      end
     end
   end
 end