From 4145aa6bcd69552ecf76b04721c57811895359c8 Mon Sep 17 00:00:00 2001
From: Mitchell Hashimoto <mitchell.hashimoto@gmail.com>
Date: Sat, 12 Apr 2014 08:45:04 -0700
Subject: [PATCH] privisioners/shell: wait for reboot if we can

---
 plugins/guests/windows/cap/reboot.rb          | 19 +++++++++
 plugins/guests/windows/plugin.rb              |  5 +++
 .../guests/windows/scripts/reboot_detect.ps1  | 42 +++++++++++++++++++
 plugins/provisioners/shell/provisioner.rb     |  4 +-
 4 files changed, 69 insertions(+), 1 deletion(-)
 create mode 100644 plugins/guests/windows/cap/reboot.rb
 create mode 100644 plugins/guests/windows/scripts/reboot_detect.ps1

diff --git a/plugins/guests/windows/cap/reboot.rb b/plugins/guests/windows/cap/reboot.rb
new file mode 100644
index 000000000..e09037cf2
--- /dev/null
+++ b/plugins/guests/windows/cap/reboot.rb
@@ -0,0 +1,19 @@
+module VagrantPlugins
+  module GuestWindows
+    module Cap
+      class Reboot
+        def self.wait_for_reboot(machine)
+          # Technically it should be possible to make it work with SSH
+          # too, but we don't yet.
+          return if machine.config.vm.communicator != :winrm
+
+          script  = File.expand_path("../../scripts/reboot_detect.ps1", __FILE__)
+          script  = File.read(script)
+          while machine.communicate.execute(script, error_check: false) != 0
+            sleep 10
+          end
+        end
+      end
+    end
+  end
+end
diff --git a/plugins/guests/windows/plugin.rb b/plugins/guests/windows/plugin.rb
index eed04d872..fa1408314 100644
--- a/plugins/guests/windows/plugin.rb
+++ b/plugins/guests/windows/plugin.rb
@@ -49,6 +49,11 @@ module VagrantPlugins
         Cap::MountSharedFolder
       end
 
+      guest_capability(:windows, :wait_for_reboot) do
+        require_relative "cap/reboot"
+        Cap::Reboot
+      end
+
       protected
 
       def self.init!
diff --git a/plugins/guests/windows/scripts/reboot_detect.ps1 b/plugins/guests/windows/scripts/reboot_detect.ps1
new file mode 100644
index 000000000..fa2e84ffe
--- /dev/null
+++ b/plugins/guests/windows/scripts/reboot_detect.ps1
@@ -0,0 +1,42 @@
+# Function to check whether machine is currently shutting down
+function ShuttingDown {
+    [string]$sourceCode = @"
+using System;
+using System.Runtime.InteropServices;
+
+namespace Vagrant {
+    public static class RemoteManager {
+        private const int SM_SHUTTINGDOWN = 0x2000;
+
+        [DllImport("User32.dll", CharSet = CharSet.Unicode)]
+        private static extern int GetSystemMetrics(int Index);
+
+        public static bool Shutdown() {
+            return (0 != GetSystemMetrics(SM_SHUTTINGDOWN));
+        }
+    }
+}
+"@
+    $type = Add-Type -TypeDefinition $sourceCode -PassThru
+    return $type::Shutdown()
+}
+
+if (ShuttingDown) {
+  exit 1
+} else {
+  # See if a reboot is scheduled in the future by trying to schedule a reboot
+  . shutdown.exe -f -r -t 60
+
+  if ($LASTEXITCODE -eq 1190) {
+    # reboot is already pending
+    exit 2
+  }
+
+  # Remove the pending reboot we just created above
+  if ($LASTEXITCODE -eq 0) {
+    . shutdown.exe -a
+  }
+}
+
+# no reboot in progress or scheduled
+exit 0
diff --git a/plugins/provisioners/shell/provisioner.rb b/plugins/provisioners/shell/provisioner.rb
index a04c51382..109826cf4 100644
--- a/plugins/provisioners/shell/provisioner.rb
+++ b/plugins/provisioners/shell/provisioner.rb
@@ -74,7 +74,9 @@ module VagrantPlugins
       # This provisions using WinRM, which assumes a PowerShell
       # console on the other side.
       def provision_winrm(args)
-        # TODO: Wait for rebooting
+        if @machine.guest.capability?(:wait_for_reboot)
+          @machine.guest.capability(:wait_for_reboot)
+        end
 
         with_script_file do |path|
           @machine.communicate.tap do |comm|