From 8c975e5c4662f3fe84348a3e9a3c0d6050a882a1 Mon Sep 17 00:00:00 2001
From: "Joshua M. Boniface" <joshua@boniface.me>
Date: Sun, 11 Jul 2021 23:10:41 -0400
Subject: [PATCH] Add chroot context manager example to debootstrap

Closes #132
---
 .../examples/debootstrap_script.py            | 80 ++++++++++---------
 .../provisioner/examples/dummy_script.py      |  2 +-
 2 files changed, 43 insertions(+), 39 deletions(-)

diff --git a/api-daemon/provisioner/examples/debootstrap_script.py b/api-daemon/provisioner/examples/debootstrap_script.py
index ed569d92..c554ce5e 100644
--- a/api-daemon/provisioner/examples/debootstrap_script.py
+++ b/api-daemon/provisioner/examples/debootstrap_script.py
@@ -34,6 +34,29 @@
 # with that.
 
 import os
+from contextlib import contextmanager
+
+
+# Create a chroot context manager
+# This can be used later in the script to chroot to the destination directory
+# for instance to run commands within the target.
+@contextmanager
+def chroot_target(destination):
+    try:
+        real_root = os.open("/", os.O_RDONLY)
+        os.chroot(destination)
+        fake_root = os.open("/", os.O_RDONLY)
+        os.fchdir(fake_root)
+        yield
+    finally:
+        os.fchdir(real_root)
+        os.chroot(".")
+        os.fchdir(real_root)
+        os.close(fake_root)
+        os.close(real_root)
+        del fake_root
+        del real_root
+
 
 # Installation function - performs a debootstrap install of a Debian system
 # Note that the only arguments are keyword arguments.
@@ -193,40 +216,25 @@ GRUB_DISABLE_LINUX_UUID=false
         fh.write(data)
 
     # Chroot, do some in-root tasks, then exit the chroot
-    # EXITING THE CHROOT IS VERY IMPORTANT OR THE FOLLOWING STAGES OF THE PROVISIONER
-    # WILL FAIL IN UNEXPECTED WAYS! Keep this in mind when using chroot in your scripts.
-    real_root = os.open("/", os.O_RDONLY)
-    os.chroot(temporary_directory)
-    fake_root = os.open("/", os.O_RDONLY)
-    os.fchdir(fake_root)
-
-    # Install and update GRUB
-    os.system(
-        "grub-install --force /dev/rbd/{}/{}_{}".format(root_disk['pool'], vm_name, root_disk['disk_id'])
-    )
-    os.system(
-        "update-grub"
-    )
-    # Set a really dumb root password [TEMPORARY]
-    os.system(
-        "echo root:test123 | chpasswd"
-    )
-    # Enable cloud-init target on (first) boot
-    # NOTE: Your user-data should handle this and disable it once done, or things get messy.
-    #       That cloud-init won't run without this hack seems like a bug... but even the official
-    #       Debian cloud images are affected, so who knows.
-    os.system(
-        "systemctl enable cloud-init.target"
-    )
-
-    # Restore our original root/exit the chroot
-    # EXITING THE CHROOT IS VERY IMPORTANT OR THE FOLLOWING STAGES OF THE PROVISIONER
-    # WILL FAIL IN UNEXPECTED WAYS! Keep this in mind when using chroot in your scripts.
-    os.fchdir(real_root)
-    os.chroot(".")
-    os.fchdir(real_root)
-    os.close(fake_root)
-    os.close(real_root)
+    with chroot_target(temporary_directory):
+        # Install and update GRUB
+        os.system(
+            "grub-install --force /dev/rbd/{}/{}_{}".format(root_disk['pool'], vm_name, root_disk['disk_id'])
+        )
+        os.system(
+            "update-grub"
+        )
+        # Set a really dumb root password [TEMPORARY]
+        os.system(
+            "echo root:test123 | chpasswd"
+        )
+        # Enable cloud-init target on (first) boot
+        # NOTE: Your user-data should handle this and disable it once done, or things get messy.
+        #       That cloud-init won't run without this hack seems like a bug... but even the official
+        #       Debian cloud images are affected, so who knows.
+        os.system(
+            "systemctl enable cloud-init.target"
+        )
 
     # Unmount the bound devfs
     os.system(
@@ -235,8 +243,4 @@ GRUB_DISABLE_LINUX_UUID=false
         )
     )
 
-    # Clean up file handles so paths can be unmounted
-    del fake_root
-    del real_root
-
     # Everything else is done via cloud-init user-data
diff --git a/api-daemon/provisioner/examples/dummy_script.py b/api-daemon/provisioner/examples/dummy_script.py
index 216b5a04..bb7bebb9 100644
--- a/api-daemon/provisioner/examples/dummy_script.py
+++ b/api-daemon/provisioner/examples/dummy_script.py
@@ -29,7 +29,7 @@
 # This script will run under root privileges as the provisioner does. Be careful
 # with that.
 
-# Installation function - performs a debootstrap install of a Debian system
+# Installation function - performs no actions then returns
 # Note that the only arguments are keyword arguments.
 def install(**kwargs):
     # The provisioner has already mounted the disks on kwargs['temporary_directory'].