From 7f1beb1d6449978a81acf61426eb47b45a7259be Mon Sep 17 00:00:00 2001
From: Michael Hudson-Doyle <michael.hudson@canonical.com>
Date: Thu, 19 Aug 2021 14:28:38 +1200
Subject: [PATCH] add a mode argument to shutdown to allow reboot or power off

---
 subiquity/client/controllers/progress.py |  7 +++++--
 subiquity/common/apidef.py               |  3 ++-
 subiquity/common/types.py                |  5 +++++
 subiquity/server/controllers/shutdown.py | 12 +++++++++---
 subiquitycore/context.py                 |  2 +-
 5 files changed, 22 insertions(+), 7 deletions(-)

diff --git a/subiquity/client/controllers/progress.py b/subiquity/client/controllers/progress.py
index a173a0f5..fba00777 100644
--- a/subiquity/client/controllers/progress.py
+++ b/subiquity/client/controllers/progress.py
@@ -21,7 +21,10 @@ import aiohttp
 from subiquitycore.context import with_context
 
 from subiquity.client.controller import SubiquityTuiController
-from subiquity.common.types import ApplicationState
+from subiquity.common.types import (
+    ApplicationState,
+    ShutdownMode,
+    )
 from subiquity.ui.views.installprogress import (
     InstallRunning,
     ProgressView,
@@ -65,7 +68,7 @@ class ProgressController(SubiquityTuiController):
 
     async def send_reboot_and_wait(self):
         try:
-            await self.app.client.shutdown.POST()
+            await self.app.client.shutdown.POST(mode=ShutdownMode.REBOOT)
         except aiohttp.ClientError:
             pass
         self.app.exit()
diff --git a/subiquity/common/apidef.py b/subiquity/common/apidef.py
index e99bbe87..3bacac94 100644
--- a/subiquity/common/apidef.py
+++ b/subiquity/common/apidef.py
@@ -37,6 +37,7 @@ from subiquity.common.types import (
     IdentityData,
     NetworkStatus,
     RefreshStatus,
+    ShutdownMode,
     SnapInfo,
     SnapListResponse,
     SnapSelection,
@@ -236,7 +237,7 @@ class API:
         def POST(tz: str): ...
 
     class shutdown:
-        def POST(immediate: bool = False): ...
+        def POST(mode: ShutdownMode, immediate: bool = False): ...
 
 
 class LinkAction(enum.Enum):
diff --git a/subiquity/common/types.py b/subiquity/common/types.py
index 6bbc6250..8461b11d 100644
--- a/subiquity/common/types.py
+++ b/subiquity/common/types.py
@@ -330,3 +330,8 @@ class SnapListResponse:
 class TimeZoneInfo:
     timezone: str
     from_geoip: bool
+
+
+class ShutdownMode(enum.Enum):
+    REBOOT = enum.auto()
+    POWEROFF = enum.auto()
diff --git a/subiquity/server/controllers/shutdown.py b/subiquity/server/controllers/shutdown.py
index fb8ccf4e..bf74f4cb 100644
--- a/subiquity/server/controllers/shutdown.py
+++ b/subiquity/server/controllers/shutdown.py
@@ -23,6 +23,7 @@ from subiquitycore.context import with_context
 from subiquitycore.utils import arun_command, run_command
 
 from subiquity.common.apidef import API
+from subiquity.common.types import ShutdownMode
 from subiquity.server.controller import SubiquityController
 from subiquity.server.controllers.install import ApplicationState
 
@@ -43,8 +44,10 @@ class ShutdownController(SubiquityController):
         self.user_shutdown_event = asyncio.Event()
         self.server_reboot_event = asyncio.Event()
         self.shuttingdown_event = asyncio.Event()
+        self.mode = ShutdownMode.REBOOT
 
-    async def POST(self, immediate: bool = False):
+    async def POST(self, mode: ShutdownMode, immediate: bool = False):
+        self.mode = mode
         self.app.controllers.Install.stop_uu()
         self.user_shutdown_event.set()
         if immediate:
@@ -92,7 +95,7 @@ class ShutdownController(SubiquityController):
         except Exception:
             log.exception("saving journal failed")
 
-    @with_context()
+    @with_context(description='mode={self.mode.name}')
     def shutdown(self, context):
         self.shuttingdown_event.set()
         if self.opts.dry_run:
@@ -101,4 +104,7 @@ class ShutdownController(SubiquityController):
             if self.app.state == ApplicationState.DONE:
                 if platform.machine() == 's390x':
                     run_command(["chreipl", "/target/boot"])
-            run_command(["/sbin/reboot"])
+            if self.mode == ShutdownMode.REBOOT:
+                run_command(["/sbin/reboot"])
+            elif self.mode == ShutdownMode.POWEROFF:
+                run_command(["/sbin/poweroff"])
diff --git a/subiquitycore/context.py b/subiquitycore/context.py
index 0ed9c399..954ec041 100644
--- a/subiquitycore/context.py
+++ b/subiquitycore/context.py
@@ -131,7 +131,7 @@ def with_context(name=None, description="", **context_kw):
                 context = self.context
             kw['context'] = context.child(
                 name=name.format(**kw),
-                description=description.format(**kw),
+                description=description.format(self=self, **kw),
                 **context_kw)
             return kw