Add public key capability to Windows guests for winssh communicator

plugins/guests/windows/cap/public_key.rb

@@ -0,0 +1,106 @@
+require "tempfile"
+
+module VagrantPlugins
+  module GuestWindows
+    module Cap
+      class PublicKey
+        def self.insert_public_key(machine, contents)
+          if machine.communicate.is_a?(CommunicatorWinSSH::Communicator)
+            winssh_insert_public_key(machine, contents)
+          else
+            raise Vagrant::Errors::SSHInsertKeyUnsupported
+          end
+        end
+
+        def self.remove_public_key(machine, contents)
+          if machine.communicate.is_a?(CommunicatorWinSSH::Communicator)
+            winssh_remove_public_key(machine, contents)
+          else
+            raise Vagrant::Errors::SSHInsertKeyUnsupported
+          end
+        end
+
+        def self.winssh_insert_public_key(machine, contents)
+          comm = machine.communicate
+          contents = contents.strip
+
+          directories = fetch_guest_paths(comm)
+          home_dir = directories[:home]
+          temp_dir = directories[:temp]
+
+          remote_ssh_dir = "#{home_dir}\\.ssh"
+          remote_upload_path = "#{temp_dir}\\vagrant-insert-pubkey-#{Time.now.to_i}"
+          remote_authkeys_path = "#{remote_ssh_dir}\authorized_keys"
+
+          # Ensure the user's ssh directory exists
+          comm.execute("dir \"#{remote_ssh_dir}\"\n if errorlevel 1 (mkdir \"#{remote_ssh_dir}\")", shell: "cmd")
+          remote_upload_path = "#{temp_dir}\\vagrant-insert-pubkey-#{Time.now.to_i}"
+          remote_authkeys_path = "#{remote_ssh_dir}\\authorized_keys"
+
+          keys_file = Tempfile.new("vagrant-windows-insert-public-key")
+          # Check if an authorized_keys file already exists
+          result = comm.execute("dir \"#{remote_authkeys_path}\"", shell: "cmd", error_check: false)
+          if result == 0
+            keys_file.close
+            comm.download(remote_authkeys_path, keys_file.path)
+            current_content = File.read(keys_file.path).split(/[\r\n]+/)
+            if !current_content.include?(contents)
+              current_content << contents
+            end
+            File.write(keys_file.path, current_content.join("\r\n") + "\r\n")
+          else
+            keys_file.puts(contents)
+            keys_file.close
+          end
+          keys_file.delete
+          comm.upload(keys_file.path, remote_upload_path)
+          comm.execute("move /y \"#{remote_upload_path}\" \"#{remote_authkeys_path}\"", shell: "cmd")
+        end
+
+        def self.winssh_remove_public_key(machine, contents)
+          comm = machine.communicate
+
+          directories = fetch_guest_paths(comm)
+          home_dir = directories[:home]
+          temp_dir = directories[:temp]
+
+          remote_ssh_dir = "#{home_dir}\\.ssh"
+          remote_upload_path = "#{temp_dir}\\vagrant-remove-pubkey-#{Time.now.to_i}"
+          remote_authkeys_path = "#{remote_ssh_dir}\\authorized_keys"
+
+          # Check if an authorized_keys file already exists
+          result = comm.execute("dir \"#{remote_authkeys_path}\"", shell: "cmd", error_check: false)
+          if result == 0
+            keys_file = Tempfile.new("vagrant-windows-remove-public-key")
+            keys_file.close
+            comm.download(remote_authkeys_path, keys_file.path)
+            current_content = File.read(keys_file.path).split(/[\r\n]+/)
+            current_content.delete(contents)
+            File.write(keys_file.path, current_content.join("\r\n") + "\r\n")
+            comm.upload(keys_file.path, remote_upload_path)
+            keys_file.delete
+            comm.execute("move /y \"#{remote_upload_path}\" \"#{remote_authkeys_path}\"", shell: "cmd")
+          end
+        end
+
+        # Fetch user's temporary and home directory paths from the Windows guest
+        #
+        # @param [Communicator]
+        # @return [Hash] {:temp, :home}
+        def self.fetch_guest_paths(communicator)
+          output = ""
+          communicator.execute("echo %TEMP%\necho %USERPROFILE%", shell: "cmd") do |type, data|
+            if type == :stdout
+              output << data
+            end
+          end
+          temp_dir, home_dir = output.strip.split(/[\r\n]+/)
+          if temp_dir.nil? || home_dir.nil?
+            raise Errors::PublicKeyDirectoryFailure
+          end
+          {temp: temp_dir, home: home_dir}
+        end
+      end
+    end
+  end
+end

plugins/guests/windows/errors.rb

@@ -13,6 +13,10 @@ module VagrantPlugins
       class RenameComputerFailed < WindowsError
         error_key(:rename_computer_failed)
       end
+
+      class PublicKeyDirectoryFailure < WindowsError
+        error_key(:public_key_directory_failure)
+      end
     end
   end
 end

plugins/guests/windows/plugin.rb

@@ -74,6 +74,16 @@ module VagrantPlugins
         Cap::RSync
       end
 
+      guest_capability(:windows, :insert_public_key) do
+        require_relative "cap/public_key"
+        Cap::PublicKey
+      end
+
+      guest_capability(:windows, :remove_public_key) do
+        require_relative "cap/public_key"
+        Cap::PublicKey
+      end
+
       protected
 
       def self.init!

templates/locales/guest_windows.yml

@@ -1,6 +1,10 @@
 en:
   vagrant_windows:
     errors:
+      public_key_directory_failure: |-
+        Vagrant failed to properly discover the correct paths for the
+        temporary directory and user profile directory on the Windows
+        guest. Please ensure the guest is properly configured.
       network_winrm_required: |-
         Configuring networks on Windows requires the communicator to be
         set to WinRM. To do this, add the following to your Vagrantfile:

test/unit/plugins/guests/windows/cap/insert_public_key_test.rb

@@ -0,0 +1,75 @@
+require "tempfile"
+require_relative "../../../../base"
+require_relative "../../../../../../plugins/communicators/winssh/communicator"
+
+describe "VagrantPlugins::GuestWindows::Cap::InsertPublicKey" do
+  let(:caps) do
+    VagrantPlugins::GuestWindows::Plugin
+      .components
+      .guest_capabilities[:windows]
+  end
+
+  let(:machine) { double("machine") }
+  let(:comm) { VagrantTests::DummyCommunicator::Communicator.new(machine) }
+  let(:auth_keys_check_result){ 1 }
+
+  before do
+    @tempfile = Tempfile.new("vagrant-test")
+    allow(Tempfile).to receive(:new).and_return(@tempfile)
+    allow(comm).to receive(:is_a?).and_return(true)
+    allow(machine).to receive(:communicate).and_return(comm)
+
+    allow(comm).to receive(:execute).with(/echo .+/, shell: "cmd").and_yield(:stdout, "TEMP\r\nHOME\r\n")
+    allow(comm).to receive(:execute).with(/dir .+\.ssh/, shell: "cmd")
+    allow(comm).to receive(:execute).with(/dir .+authorized_keys/, shell: "cmd", error_check: false).and_return(auth_keys_check_result)
+  end
+
+  after do
+    @tempfile.delete
+  end
+
+  describe ".insert_public_key" do
+    let(:cap) { caps.get(:insert_public_key) }
+
+    context "when authorized_keys exists on guest" do
+      let(:auth_keys_check_result){ 0 }
+      before do
+        expect(@tempfile).to receive(:delete).and_return(true)
+        expect(@tempfile).to receive(:delete).and_call_original
+      end
+
+      it "inserts the public key" do
+        expect(comm).to receive(:download)
+        expect(comm).to receive(:upload)
+        expect(comm).to receive(:execute).with(/move .*/, shell: "cmd")
+        cap.insert_public_key(machine, "ssh-rsa ...")
+        expect(File.read(@tempfile.path)).to include("ssh-rsa ...")
+      end
+    end
+
+    context "when authorized_keys does not exist on guest" do
+      before do
+        expect(@tempfile).to receive(:delete).and_return(true)
+        expect(@tempfile).to receive(:delete).and_call_original
+      end
+
+      it "inserts the public key" do
+        expect(comm).to_not receive(:download)
+        expect(comm).to receive(:upload)
+        expect(comm).to receive(:execute).with(/move .*/, shell: "cmd")
+        cap.insert_public_key(machine, "ssh-rsa ...")
+        expect(File.read(@tempfile.path)).to include("ssh-rsa ...")
+      end
+    end
+
+    context "when required directories cannot be fetched from the guest" do
+      before do
+        expect(comm).to receive(:execute).with(/echo .+/, shell: "cmd").and_yield(:stdout, "TEMP\r\n")
+      end
+
+      it "should raise an error" do
+        expect{ cap.insert_public_key(machine, "ssh-rsa ...") }.to raise_error(VagrantPlugins::GuestWindows::Errors::PublicKeyDirectoryFailure)
+      end
+    end
+  end
+end

test/unit/plugins/guests/windows/cap/remove_public_key_test.rb

@@ -0,0 +1,66 @@
+require "tempfile"
+require_relative "../../../../base"
+require_relative "../../../../../../plugins/communicators/winssh/communicator"
+
+describe "VagrantPlugins::GuestWindows::Cap::RemovePublicKey" do
+  let(:caps) do
+    VagrantPlugins::GuestWindows::Plugin
+      .components
+      .guest_capabilities[:windows]
+  end
+
+  let(:machine) { double("machine") }
+  let(:comm) { VagrantTests::DummyCommunicator::Communicator.new(machine) }
+  let(:public_key_insecure){ "ssh-rsa...insecure" }
+  let(:public_key_other){ "ssh-rsa...other" }
+
+  let(:auth_keys_check_result){ 1 }
+
+  before do
+    @tempfile = Tempfile.new("vagrant-test")
+    @tempfile.puts(public_key_insecure)
+    @tempfile.puts(public_key_other)
+    @tempfile.flush
+    @tempfile.rewind
+    allow(Tempfile).to receive(:new).and_return(@tempfile)
+    allow(comm).to receive(:is_a?).and_return(true)
+    allow(machine).to receive(:communicate).and_return(comm)
+
+    allow(comm).to receive(:execute).with(/echo .+/, shell: "cmd").and_yield(:stdout, "TEMP\r\nHOME\r\n")
+    allow(comm).to receive(:execute).with(/dir .+authorized_keys/, shell: "cmd", error_check: false).and_return(auth_keys_check_result)
+  end
+
+  after do
+    @tempfile.delete
+  end
+
+  describe ".remove_public_key" do
+    let(:cap) { caps.get(:remove_public_key) }
+
+    context "when authorized_keys exists on guest" do
+      let(:auth_keys_check_result){ 0 }
+      before do
+        expect(@tempfile).to receive(:delete).and_return(true)
+        expect(@tempfile).to receive(:delete).and_call_original
+      end
+
+      it "removes the public key" do
+        expect(comm).to receive(:download)
+        expect(comm).to receive(:upload)
+        expect(comm).to receive(:execute).with(/move .*/, shell: "cmd")
+        cap.remove_public_key(machine, public_key_insecure)
+        expect(File.read(@tempfile.path)).to include(public_key_other)
+        expect(File.read(@tempfile.path)).to_not include(public_key_insecure)
+      end
+    end
+
+    context "when authorized_keys does not exist on guest" do
+      it "does nothing" do
+        expect(comm).to_not receive(:download)
+        expect(comm).to_not receive(:upload)
+        expect(comm).to_not receive(:execute).with(/move .*/, shell: "cmd")
+        cap.remove_public_key(machine, public_key_insecure)
+      end
+    end
+  end
+end