From cdeac10b3622be9d932095e3b25ccb354f15ebd1 Mon Sep 17 00:00:00 2001
From: Tobias <tsmolka@gmail.com>
Date: Tue, 28 Feb 2017 02:35:00 -0800
Subject: [PATCH 1/2] guests/esxi: Added public_key capability

---
 plugins/guests/esxi/cap/public_key.rb | 63 +++++++++++++++++++++++++++
 plugins/guests/esxi/plugin.rb         | 11 +++++
 2 files changed, 74 insertions(+)
 create mode 100644 plugins/guests/esxi/cap/public_key.rb

diff --git a/plugins/guests/esxi/cap/public_key.rb b/plugins/guests/esxi/cap/public_key.rb
new file mode 100644
index 000000000..84f199c39
--- /dev/null
+++ b/plugins/guests/esxi/cap/public_key.rb
@@ -0,0 +1,63 @@
+require "tempfile"
+
+require "vagrant/util/shell_quote"
+
+module VagrantPlugins
+  module GuestEsxi
+    module Cap
+      class PublicKey
+      
+        def self.insert_public_key(machine, contents)
+          comm = machine.communicate
+          contents = contents.strip << "\n"
+
+          remote_path = "/tmp/vagrant-insert-pubkey-#{Time.now.to_i}"
+          Tempfile.open("vagrant-esxi-insert-public-key") do |f|
+            f.binmode
+            f.write(contents)
+            f.fsync
+            f.close
+            comm.upload(f.path, remote_path)
+          end
+
+          comm.execute <<-EOH.gsub(/^ {12}/, "")
+            set -e
+            SSH_DIR="$(grep -q '^AuthorizedKeysFile\s*\/etc\/ssh\/keys-%u\/authorized_keys$' /etc/ssh/sshd_config && echo -n /etc/ssh/keys-${USER} || echo -n ~/.ssh)"
+            mkdir -p "${SSH_DIR}"
+            chmod 0700 "${SSH_DIR}"
+            cat '#{remote_path}' >> "${SSH_DIR}/authorized_keys"
+            chmod 0600 "${SSH_DIR}/authorized_keys"
+            rm -f '#{remote_path}'
+          EOH
+        end
+
+        def self.remove_public_key(machine, contents)
+          comm = machine.communicate
+          contents = contents.strip << "\n"
+
+          remote_path = "/tmp/vagrant-remove-pubkey-#{Time.now.to_i}"
+          Tempfile.open("vagrant-esxi-remove-public-key") do |f|
+            f.binmode
+            f.write(contents)
+            f.fsync
+            f.close
+            comm.upload(f.path, remote_path)
+          end
+
+          # Use execute (not sudo) because we want to execute this as the SSH
+          # user (which is "vagrant" by default).
+          comm.execute <<-EOH.sub(/^ {12}/, "")
+            set -e
+            SSH_DIR="$(grep -q '^AuthorizedKeysFile\s*\/etc\/ssh\/keys-%u\/authorized_keys$' /etc/ssh/sshd_config && echo -n /etc/ssh/keys-${USER} || echo -n ~/.ssh)"
+            if test -f "${SSH_DIR}/authorized_keys"; then
+              grep -v -x -f '#{remote_path}' "${SSH_DIR}/authorized_keys" > "${SSH_DIR}/authorized_keys.tmp"
+              mv "${SSH_DIR}/authorized_keys.tmp" "${SSH_DIR}/authorized_keys"
+              chmod 0600 "${SSH_DIR}/authorized_keys"
+            fi
+            rm -f '#{remote_path}'
+          EOH
+        end
+      end
+    end
+  end
+end
diff --git a/plugins/guests/esxi/plugin.rb b/plugins/guests/esxi/plugin.rb
index 04367bc22..8d68e1873 100644
--- a/plugins/guests/esxi/plugin.rb
+++ b/plugins/guests/esxi/plugin.rb
@@ -30,6 +30,17 @@ module VagrantPlugins
         require_relative "cap/halt"
         Cap::Halt
       end
+      
+      guest_capability(:esxi, :remove_public_key) do
+        require_relative "cap/public_key"
+        Cap::PublicKey
+      end
+      
+      guest_capability(:esxi, :insert_public_key) do
+        require_relative "cap/public_key"
+        Cap::PublicKey
+      end
+      
     end
   end
 end

From d074ac7906f2068b96a4987190d91ef62ff82f16 Mon Sep 17 00:00:00 2001
From: Tobias <tsmolka@gmail.com>
Date: Tue, 28 Feb 2017 08:22:05 -0800
Subject: [PATCH 2/2] guests/esxi: Added public_key_test

---
 .../guests/esxi/cap/public_key_test.rb        | 48 +++++++++++++++++++
 1 file changed, 48 insertions(+)
 create mode 100644 test/unit/plugins/guests/esxi/cap/public_key_test.rb

diff --git a/test/unit/plugins/guests/esxi/cap/public_key_test.rb b/test/unit/plugins/guests/esxi/cap/public_key_test.rb
new file mode 100644
index 000000000..93a300852
--- /dev/null
+++ b/test/unit/plugins/guests/esxi/cap/public_key_test.rb
@@ -0,0 +1,48 @@
+require_relative "../../../../base"
+
+describe "VagrantPlugins::GuestEsxi::Cap::PublicKey" do
+  let(:caps) do
+    VagrantPlugins::GuestEsxi::Plugin
+      .components
+      .guest_capabilities[:esxi]
+  end
+
+  let(:machine) { double("machine") }
+  let(:comm) { VagrantTests::DummyCommunicator::Communicator.new(machine) }
+
+  before do
+    allow(machine).to receive(:communicate).and_return(comm)
+  end
+
+  after do
+    comm.verify_expectations!
+  end
+
+  describe ".insert_public_key" do
+    let(:cap) { caps.get(:insert_public_key) }
+
+    it "inserts the public key" do
+      cap.insert_public_key(machine, "ssh-rsa ...")
+      expect(comm.received_commands[0]).to match(/SSH_DIR=".*"/)
+      expect(comm.received_commands[0]).to match(/mkdir -p "\${SSH_DIR}"/)
+      expect(comm.received_commands[0]).to match(/chmod 0700 "\${SSH_DIR}"/)
+      expect(comm.received_commands[0]).to match(/cat '\/tmp\/vagrant-(.+)' >> "\${SSH_DIR}\/authorized_keys"/)
+      expect(comm.received_commands[0]).to match(/chmod 0600 "\${SSH_DIR}\/authorized_keys"/)
+      expect(comm.received_commands[0]).to match(/rm -f '\/tmp\/vagrant-(.+)'/)
+    end
+  end
+
+  describe ".remove_public_key" do
+    let(:cap) { caps.get(:remove_public_key) }
+
+    it "removes the public key" do
+      cap.remove_public_key(machine, "ssh-rsa ...")
+      expect(comm.received_commands[0]).to match(/SSH_DIR=".*"/)
+      expect(comm.received_commands[0]).to match(/grep -v -x -f '\/tmp\/vagrant-(.+)' "\${SSH_DIR}\/authorized_keys" > "\${SSH_DIR}\/authorized_keys\.tmp"/)
+      expect(comm.received_commands[0]).to match(/mv "\${SSH_DIR}\/authorized_keys\.tmp" "\${SSH_DIR}\/authorized_keys"/)
+      expect(comm.received_commands[0]).to match(/chmod 0600 "\${SSH_DIR}\/authorized_keys"/)
+      expect(comm.received_commands[0]).to match(/rm -f '\/tmp\/vagrant-(.+)'/)
+    end
+  end
+
+end