diff --git a/usr/lib/akshara/utils/gen_rootfs.py b/usr/lib/akshara/utils/gen_rootfs.py
index c68990e..7384ee3 100644
--- a/usr/lib/akshara/utils/gen_rootfs.py
+++ b/usr/lib/akshara/utils/gen_rootfs.py
@@ -8,23 +8,28 @@ from . import output
 
 
 def run_script_rootfs(rootfs: RootFS, input: str, args: list):
-    """Run a script within the rootfs."""
+    """Runs a script within the rootfs."""
 
     rootfs.exec(["bash", "-s", *args], text=True, input=input)
 
 
-def gen_rootfs(system_config: dict, rootfs_path: str, env: dict = {}) -> RootFS:
+def gen_rootfs(system_config: dict, rootfs_path: str) -> RootFS:
     """Generates a rootfs for a given system configuration."""
 
-    rootfs = RootFS(rootfs_path, system_config["distro-config"], env)
+    rootfs = RootFS(rootfs_path, system_config["distro-config"], system_config["env"])
 
-    subprocess.run(
-        ["bash", "-s"],
-        text=True,
-        input=system_config["distro-config"]["before-stages"],
-        cwd=rootfs_path,
-        env=os.environ.copy() | env,
-    )
+    if (
+        subprocess.run(
+            ["bash", "-s"],
+            text=True,
+            input=system_config["distro-config"]["before-stages"],
+            cwd=rootfs_path,
+            env=os.environ.copy() | system_config["env"],
+        ).returncode
+        != 0
+    ):
+        output.error("failed to initialise rootfs")
+        exit(1)
 
     modules = {}
 
@@ -42,13 +47,18 @@ def gen_rootfs(system_config: dict, rootfs_path: str, env: dict = {}) -> RootFS:
             run_script_rootfs(rootfs, modules[stage["module"]], inputs)
 
     if isinstance(system_config["distro-config"].get("after-stages"), str):
-        subprocess.run(
-            ["bash", "-s"],
-            text=True,
-            input=system_config["distro-config"]["after-stages"],
-            cwd=str(rootfs_path),
-            env=os.environ.copy() | env,
-        )
+        if (
+            subprocess.run(
+                ["bash", "-s"],
+                text=True,
+                input=system_config["distro-config"]["after-stages"],
+                cwd=str(rootfs_path),
+                env=os.environ.copy() | system_config["env"],
+            )
+            != 0
+        ):
+            output.error("failed to finalise rootfs")
+            exit(1)
 
     with open(os.path.join(rootfs_path, "usr/system.json"), "w") as system_json_file:
         json.dump(system_config, system_json_file, ensure_ascii=False)
diff --git a/usr/lib/akshara/utils/helpers.py b/usr/lib/akshara/utils/helpers.py
index 874cc33..93e903f 100644
--- a/usr/lib/akshara/utils/helpers.py
+++ b/usr/lib/akshara/utils/helpers.py
@@ -119,6 +119,7 @@ def is_already_latest() -> bool:
                             ["bash", "-s"],
                             text=True,
                             input=system_config["distro-config"]["should-update"],
+                            env=os.environ.copy() | system_config["env"],
                         ).returncode
                         != 0
                     )