1 files changed, 88 insertions, 6 deletions

diff --git a/Software/Visual_Studio/Tango.Integration/Operation/MachineOperator.cs b/Software/Visual_Studio/Tango.Integration/Operation/MachineOperator.cs
index d0e03f873..d421120ea 100644
--- a/Software/Visual_Studio/Tango.Integration/Operation/MachineOperator.cs
+++ b/Software/Visual_Studio/Tango.Integration/Operation/MachineOperator.cs
@@ -33,6 +33,7 @@ using Tango.Integration.Upgrade;
 using Tango.PMR.FirmwareUpgrade;
 using Tango.Integration.Logging;
 using Tango.Integration.JobRuns;
+using Tango.FirmwareUpdateLib.WPF;
 
 namespace Tango.Integration.Operation
 {
@@ -86,6 +87,7 @@ namespace Tango.Integration.Operation
             EnableEventsNotification = true;
             EnableJobResume = true;
             LogEmbeddedDebuggingToFile = true;
+            FirmwareUpgradeMode = FirmwareUpgradeModes.DFU | FirmwareUpgradeModes.TFP_PACKAGE;
         }
 
         /// <summary>
@@ -202,6 +204,11 @@ namespace Tango.Integration.Operation
         }
 
         /// <summary>
+        /// Gets or sets the firmware upgrade mode.
+        /// </summary>
+        public FirmwareUpgradeModes FirmwareUpgradeMode { get; set; }
+
+        /// <summary>
         /// Gets a value indicating whether this instance is printing.
         /// </summary>
         public bool IsPrinting
@@ -677,6 +684,8 @@ namespace Tango.Integration.Operation
         /// <returns></returns>
         public async override Task Disconnect()
         {
+            if (Status == MachineStatuses.Upgrading) return;
+
             if (State == TransportComponentState.Connected)
             {
                 DisconnectRequest request = new DisconnectRequest();
@@ -711,7 +720,11 @@ namespace Tango.Integration.Operation
         {
             var keep_alive = UseKeepAlive;
             UseKeepAlive = false;
-            await base.Connect();
+
+            if (Status != MachineStatuses.Upgrading)
+            {
+                await base.Connect();
+            }
 
             if (State == TransportComponentState.Connected)
             {
@@ -723,7 +736,10 @@ namespace Tango.Integration.Operation
                     var response = await SendRequest<ConnectRequest, ConnectResponse>(request);
                     LogResponseReceived(response.Message);
 
-                    Status = MachineStatuses.ReadyToDye;
+                    if (Status != MachineStatuses.Upgrading)
+                    {
+                        Status = MachineStatuses.ReadyToDye;
+                    }
 
                     DeviceInformation = response.Message.DeviceInformation;
 
@@ -1280,7 +1296,7 @@ namespace Tango.Integration.Operation
 
             request.JobTicket.UploadStrategy = JobUploadStrategy;
 
-            ThreadFactory.StartNew(async () => 
+            ThreadFactory.StartNew(async () =>
             {
                 if (JobUploadStrategy == JobUploadStrategy.JobDescriptionFile)
                 {
@@ -2023,9 +2039,17 @@ namespace Tango.Integration.Operation
 
             try
             {
-                zip = ZipFile.Read(tfpStream);
+                if (Status != MachineStatuses.ReadyToDye)
+                {
+                    throw LogManager.Log(new InvalidOperationException($"Could not perform firmware upgrade while operator status is '{Status}'."));
+                }
+
+                Status = MachineStatuses.Upgrading;
 
-                upgradeHandler.Total = zip.Entries.Sum(x => x.UncompressedSize);
+                var package_info = await GetFirmwarePackageInfo(tfpStream);
+                tfpStream.Position = 0;
+
+                zip = ZipFile.Read(tfpStream);
 
                 var storage = CreateStorageManager();
                 var drive = await storage.GetStorageDrive();
@@ -2044,11 +2068,67 @@ namespace Tango.Integration.Operation
                 List<ZipEntry> entries = zip.Entries.ToList();
                 List<Stream> streams = new List<Stream>();
 
+                var keepAlive = UseKeepAlive;
+                UseKeepAlive = false;
+
+                Action upgradeDFU = null;
                 Action uploadNext = null;
                 Action validate = null;
                 Action activate = null;
                 Action postActivation = null;
 
+                upgradeDFU = new Action(() =>
+                {
+                    try
+                    {
+                        if (FirmwareUpgradeMode.HasFlag(FirmwareUpgradeModes.DFU))
+                        {
+                            var mcuEntry = zip.Entries.Single(x => x.FileName == package_info.FileDescriptors.Single(y => y.Destination == VersionFileDestination.Mcu).FileName);
+                            MemoryStream ms = new MemoryStream();
+                            mcuEntry.Extract(ms);
+                            ms.Position = 0;
+                            byte[] data = ms.ToArray();
+                            ms.Dispose();
+
+                            FirmwareUpgradeManager upgradeManager = new FirmwareUpgradeManager();
+                            upgradeManager.UpgradeProgress += (sender, e) =>
+                            {
+                                upgradeHandler.Total = 100;
+                                upgradeHandler.RaiseProgress((long)e.Progress, FirmwareUpgradeStatus.Upgrading, e.State.ToDescription());
+                            };
+
+                            Adapter.Disconnect().Wait();
+
+                            if (MachineEventsStateProvider != null)
+                            {
+                                MachineEventsStateProvider.Reset();
+                            }
+
+                            upgradeManager.PerformUpgrade(data).Wait();
+                            upgradeHandler.Total = zip.Entries.Sum(x => x.UncompressedSize);
+
+                            Thread.Sleep(2000);
+                            Adapter.Connect().Wait();
+                            Connect().Wait();
+                            Status = MachineStatuses.Upgrading;
+                        }
+
+                        if (FirmwareUpgradeMode.HasFlag(FirmwareUpgradeModes.TFP_PACKAGE))
+                        {
+                            uploadNext();
+                        }
+                        else
+                        {
+                            postActivation();
+                        }
+                    }
+                    catch (Exception ex)
+                    {
+                        upgradeHandler.RaiseFailed(ex);
+                        return;
+                    }
+                });
+
                 uploadNext = new Action(() =>
                 {
                     if (entries.Count > 0)
@@ -2124,11 +2204,13 @@ namespace Tango.Integration.Operation
                 postActivation = new Action(() =>
                 {
                     upgradeHandler.RaiseCompleted();
+                    Status = MachineStatuses.ReadyToDye;
+                    UseKeepAlive = keepAlive;
                 });
 
                 ThreadFactory.StartNew(() =>
                 {
-                    uploadNext();
+                    upgradeDFU();
                 });
 
                 return upgradeHandler;