From ead42bc86044d92b8261fda8719a5bff28d0db2f Mon Sep 17 00:00:00 2001
From: "Carl D. Roth" <carl.roth@bigswitch.com>
Date: Mon, 24 Oct 2016 13:11:21 -0700
Subject: [PATCH] Added upgrade-shell

---
 .../vendor-config-onl/src/bin/upgrade-shell   |  7 +++
 .../src/python/onl/install/ShellApp.py        | 48 ++++++++++++++++++-
 2 files changed, 54 insertions(+), 1 deletion(-)
 create mode 100755 packages/base/all/vendor-config-onl/src/bin/upgrade-shell

diff --git a/packages/base/all/vendor-config-onl/src/bin/upgrade-shell b/packages/base/all/vendor-config-onl/src/bin/upgrade-shell
new file mode 100755
index 00000000..7e488d28
--- /dev/null
+++ b/packages/base/all/vendor-config-onl/src/bin/upgrade-shell
@@ -0,0 +1,7 @@
+#!/usr/bin/python
+
+"""Run the upgrade image
+"""
+
+import onl.install.ShellApp
+onl.install.ShellApp.Upgrader.main()
diff --git a/packages/base/all/vendor-config-onl/src/python/onl/install/ShellApp.py b/packages/base/all/vendor-config-onl/src/python/onl/install/ShellApp.py
index 19cfaf6f..79c40e59 100644
--- a/packages/base/all/vendor-config-onl/src/python/onl/install/ShellApp.py
+++ b/packages/base/all/vendor-config-onl/src/python/onl/install/ShellApp.py
@@ -16,8 +16,9 @@ from InstallUtils import BlkidParser
 import Fit
 
 import onl.platform.current
+from onl.sysconfig import sysconfig
 
-class AppBase(SubprocessMixin):
+class AppBase(SubprocessMixin, object):
 
     @property
     def PROG(self):
@@ -109,6 +110,7 @@ class AppBase(SubprocessMixin):
         sys.exit(code)
 
 class Onie(AppBase):
+    """Application shell that uses the ONIE runtime."""
 
     PROG = "onie-shell"
 
@@ -193,6 +195,7 @@ class Onie(AppBase):
         return 1
 
 class Loader(AppBase):
+    """Application shell that uses the (installed) loader runtime."""
 
     PROG = "loader-shell"
 
@@ -309,6 +312,49 @@ class Loader(AppBase):
         self.log.error("invalid platform-config")
         return 1
 
+class Upgrader(AppBase):
+    """Application shell that uses on-disk upgrade loader runtime."""
+
+    PROG = "upgrade-shell"
+
+    def runGrub(self):
+
+        d = sysconfig.upgrade.loader.package.dir
+        for b in sysconfig.upgrade.loader.package.grub:
+            p = os.path.join(d, b)
+            if os.path.exists(p):
+                self.log.debug("found upgrade initrd at %s", p)
+                return self._runInitrdShell(p)
+
+        self.log.error("cannot find upgrade initrd")
+        return 1
+
+    def runUboot(self):
+
+        d = sysconfig.upgrade.loader.package.dir
+        for b in sysconfig.upgrade.loader.package.fit:
+            p = os.path.join(d, b)
+            if os.path.exists(p):
+                self.log.debug("found upgrade FIT image %s", p)
+                return self._runFitShell(p)
+
+        self.log.error("cannot find FIT image")
+        return 1
+
+    def run(self):
+
+        self.platform = onl.platform.current.OnlPlatform()
+        self.pc = self.platform.platform_config
+
+        if 'grub' in self.pc:
+            return self.runGrub()
+
+        if 'flat_image_tree' in self.pc:
+            return self.runUboot()
+
+        self.log.error("invalid platform-config")
+        return 1
+
 main = Onie.main
 
 if __name__ == "__main__":