From 4aaa600bd6c677b5476ab47ed1cd1d8a04b35a95 Mon Sep 17 00:00:00 2001
From: Seth Vargo <sethvargo@gmail.com>
Date: Fri, 24 Jun 2016 21:19:08 -0400
Subject: [PATCH] guests/linux: Centralize public_key logic

---
 plugins/guests/linux/cap/insert_public_key.rb | 31 ---------
 plugins/guests/linux/cap/public_key.rb        | 65 +++++++++++++++++++
 plugins/guests/linux/cap/remove_public_key.rb | 24 -------
 plugins/guests/linux/plugin.rb                |  8 +--
 4 files changed, 69 insertions(+), 59 deletions(-)
 delete mode 100644 plugins/guests/linux/cap/insert_public_key.rb
 create mode 100644 plugins/guests/linux/cap/public_key.rb
 delete mode 100644 plugins/guests/linux/cap/remove_public_key.rb

diff --git a/plugins/guests/linux/cap/insert_public_key.rb b/plugins/guests/linux/cap/insert_public_key.rb
deleted file mode 100644
index 653720587..000000000
--- a/plugins/guests/linux/cap/insert_public_key.rb
+++ /dev/null
@@ -1,31 +0,0 @@
-module VagrantPlugins
-  module GuestLinux
-    module Cap
-      class InsertPublicKey
-        def self.insert_public_key(machine, contents)
-          comm = machine.communicate
-          contents = contents.strip << "\n"
-
-          remote_path = "/tmp/vagrant-authorized-keys-#{Time.now.to_i}"
-          Tempfile.open("vagrant-linux-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}/, '')
-            mkdir -p ~/.ssh
-            chmod 0700 ~/.ssh
-            cat '#{remote_path}' >> ~/.ssh/authorized_keys
-            chmod 0600 ~/.ssh/authorized_keys
-
-            # Remove the temporary file
-            rm -f '#{remote_path}'
-          EOH
-        end
-      end
-    end
-  end
-end
diff --git a/plugins/guests/linux/cap/public_key.rb b/plugins/guests/linux/cap/public_key.rb
new file mode 100644
index 000000000..92a718281
--- /dev/null
+++ b/plugins/guests/linux/cap/public_key.rb
@@ -0,0 +1,65 @@
+require "tempfile"
+
+require "vagrant/util/shell_quote"
+
+module VagrantPlugins
+  module GuestLinux
+    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-linux-insert-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.gsub(/^ {12}/, "")
+            set -e
+
+            mkdir -p ~/.ssh
+            chmod 0700 ~/.ssh
+            cat '#{remote_path}' >> ~/.ssh/authorized_keys
+            chmod 0600 ~/.ssh/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-bsd-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
+
+            if test -f ~/.ssh/authorized_keys; then
+              grep -v -x -f '#{remote_path}' ~/.ssh/authorized_keys > ~/.ssh/authorized_keys.tmp
+              mv ~/.ssh/authorized_keys.tmp ~/.ssh/authorized_keys
+            fi
+
+            rm -f '#{remote_path}'
+          EOH
+        end
+      end
+    end
+  end
+end
diff --git a/plugins/guests/linux/cap/remove_public_key.rb b/plugins/guests/linux/cap/remove_public_key.rb
deleted file mode 100644
index dfd9daae5..000000000
--- a/plugins/guests/linux/cap/remove_public_key.rb
+++ /dev/null
@@ -1,24 +0,0 @@
-require "vagrant/util/shell_quote"
-
-module VagrantPlugins
-  module GuestLinux
-    module Cap
-      class RemovePublicKey
-        def self.remove_public_key(machine, contents)
-          contents = contents.chomp
-          contents = Vagrant::Util::ShellQuote.escape(contents, "'")
-
-          machine.communicate.tap do |comm|
-            if comm.test("test -f ~/.ssh/authorized_keys")
-              comm.execute(<<SCRIPT)
-sed -e '/^.*#{contents}.*$/d' ~/.ssh/authorized_keys > ~/.ssh/authorized_keys.new
-mv ~/.ssh/authorized_keys.new ~/.ssh/authorized_keys
-chmod 600 ~/.ssh/authorized_keys
-SCRIPT
-            end
-          end
-        end
-      end
-    end
-  end
-end
diff --git a/plugins/guests/linux/plugin.rb b/plugins/guests/linux/plugin.rb
index 966f56353..cb0a53c19 100644
--- a/plugins/guests/linux/plugin.rb
+++ b/plugins/guests/linux/plugin.rb
@@ -22,8 +22,8 @@ module VagrantPlugins
       end
 
       guest_capability(:linux, :insert_public_key) do
-        require_relative "cap/insert_public_key"
-        Cap::InsertPublicKey
+        require_relative "cap/public_key"
+        Cap::PublicKey
       end
 
       guest_capability(:linux, :shell_expand_guest_path) do
@@ -68,8 +68,8 @@ module VagrantPlugins
       end
 
       guest_capability(:linux, :remove_public_key) do
-        require_relative "cap/remove_public_key"
-        Cap::RemovePublicKey
+        require_relative "cap/public_key"
+        Cap::PublicKey
       end
 
       guest_capability(:linux, :rsync_installed) do