5 files changed, 33 insertions, 14 deletions

diff --git a/src/api/burn/WixToolset.Mba.Core/BootstrapperCommand.cs b/src/api/burn/WixToolset.Mba.Core/BootstrapperCommand.cs
index 345e0448..88a9b9bb 100644
--- a/src/api/burn/WixToolset.Mba.Core/BootstrapperCommand.cs
+++ b/src/api/burn/WixToolset.Mba.Core/BootstrapperCommand.cs
@@ -17,7 +17,6 @@ namespace WixToolset.Mba.Core
         /// </summary>
         /// <param name="action"></param>
         /// <param name="display"></param>
-        /// <param name="restart"></param>
         /// <param name="commandLine"></param>
         /// <param name="cmdShow"></param>
         /// <param name="resume"></param>
@@ -30,7 +29,6 @@ namespace WixToolset.Mba.Core
         public BootstrapperCommand(
             LaunchAction action,
             Display display,
-            Restart restart,
             string commandLine,
             int cmdShow,
             ResumeType resume,
@@ -43,7 +41,6 @@ namespace WixToolset.Mba.Core
         {
             this.Action = action;
             this.Display = display;
-            this.Restart = restart;
             this.CommandLine = commandLine;
             this.CmdShow = cmdShow;
             this.Resume = resume;
@@ -62,9 +59,6 @@ namespace WixToolset.Mba.Core
         public Display Display { get; }
 
         /// <inheritdoc/>
-        public Restart Restart { get; }
-
-        /// <inheritdoc/>
         public string CommandLine { get; }
 
         /// <inheritdoc/>
@@ -97,6 +91,7 @@ namespace WixToolset.Mba.Core
             var args = ParseCommandLineToArgs(this.CommandLine);
             var unknownArgs = new List<string>();
             var variables = new List<KeyValuePair<string, string>>();
+            var restart = Restart.Unknown;
 
             foreach (var arg in args)
             {
@@ -104,7 +99,25 @@ namespace WixToolset.Mba.Core
 
                 if (arg[0] == '-' || arg[0] == '/')
                 {
-                    unknownArg = true;
+                    var parameter = arg.Substring(1).ToLowerInvariant();
+                    switch (parameter)
+                    {
+                        case "norestart":
+                            if (restart == Restart.Unknown)
+                            {
+                                restart = Restart.Never;
+                            }
+                            break;
+                        case "forcerestart":
+                            if (restart == Restart.Unknown)
+                            {
+                                restart = Restart.Always;
+                            }
+                            break;
+                        default:
+                            unknownArg = true;
+                            break;
+                    }
                 }
                 else
                 {
@@ -127,8 +140,14 @@ namespace WixToolset.Mba.Core
                 }
             }
 
+            if (restart == Restart.Unknown)
+            {
+                restart = this.Display < Display.Full ? Restart.Automatic : Restart.Prompt;
+            }
+
             return new MbaCommand
             {
+                Restart = restart,
                 UnknownCommandLineArgs = unknownArgs.ToArray(),
                 Variables = variables.ToArray(),
             };
diff --git a/src/api/burn/WixToolset.Mba.Core/IBootstrapperApplicationFactory.cs b/src/api/burn/WixToolset.Mba.Core/IBootstrapperApplicationFactory.cs
index 0f9193d0..64e25ff4 100644
--- a/src/api/burn/WixToolset.Mba.Core/IBootstrapperApplicationFactory.cs
+++ b/src/api/burn/WixToolset.Mba.Core/IBootstrapperApplicationFactory.cs
@@ -35,7 +35,6 @@ namespace WixToolset.Mba.Core
         [MarshalAs(UnmanagedType.I4)] internal int cbSize;
         [MarshalAs(UnmanagedType.U4)] private readonly LaunchAction action;
         [MarshalAs(UnmanagedType.U4)] private readonly Display display;
-        [MarshalAs(UnmanagedType.U4)] private readonly Restart restart;
         private readonly IntPtr wzCommandLine;
         [MarshalAs(UnmanagedType.I4)] private readonly int nCmdShow;
         [MarshalAs(UnmanagedType.U4)] private readonly ResumeType resume;
@@ -51,7 +50,6 @@ namespace WixToolset.Mba.Core
             return new BootstrapperCommand(
                 this.action,
                 this.display,
-                this.restart,
                 Marshal.PtrToStringUni(this.wzCommandLine),
                 this.nCmdShow,
                 this.resume,
diff --git a/src/api/burn/WixToolset.Mba.Core/IBootstrapperCommand.cs b/src/api/burn/WixToolset.Mba.Core/IBootstrapperCommand.cs
index b7baa55c..f583e619 100644
--- a/src/api/burn/WixToolset.Mba.Core/IBootstrapperCommand.cs
+++ b/src/api/burn/WixToolset.Mba.Core/IBootstrapperCommand.cs
@@ -20,11 +20,6 @@ namespace WixToolset.Mba.Core
         Display Display { get; }
 
         /// <summary>
-        /// Gets the action to perform if a reboot is required.
-        /// </summary>
-        Restart Restart { get; }
-
-        /// <summary>
         /// Gets the command line arguments.
         /// </summary>
         /// <returns>
diff --git a/src/api/burn/WixToolset.Mba.Core/IMbaCommand.cs b/src/api/burn/WixToolset.Mba.Core/IMbaCommand.cs
index a3ad68d8..495b2f44 100644
--- a/src/api/burn/WixToolset.Mba.Core/IMbaCommand.cs
+++ b/src/api/burn/WixToolset.Mba.Core/IMbaCommand.cs
@@ -10,6 +10,11 @@ namespace WixToolset.Mba.Core
     public interface IMbaCommand
     {
         /// <summary>
+        /// Gets the action to perform if a reboot is required.
+        /// </summary>
+        Restart Restart { get; }
+
+        /// <summary>
         /// The command line arguments not parsed into <see cref="IBootstrapperCommand"/> or <see cref="IMbaCommand"/>.
         /// </summary>
         string[] UnknownCommandLineArgs { get; }
diff --git a/src/api/burn/WixToolset.Mba.Core/MbaCommand.cs b/src/api/burn/WixToolset.Mba.Core/MbaCommand.cs
index 424cde63..8917a334 100644
--- a/src/api/burn/WixToolset.Mba.Core/MbaCommand.cs
+++ b/src/api/burn/WixToolset.Mba.Core/MbaCommand.cs
@@ -9,6 +9,8 @@ namespace WixToolset.Mba.Core
     /// </summary>
     internal sealed class MbaCommand : IMbaCommand
     {
+        public Restart Restart { get; internal set; }
+
         public string[] UnknownCommandLineArgs { get; internal set; }
 
         public KeyValuePair<string, string>[] Variables { get; internal set; }