From 8e7e5cb576d751539e5e61a2801e098db2873ffd Mon Sep 17 00:00:00 2001
From: Rudra Saraswat <rs2009@ruds.io>
Date: Thu, 2 Apr 2026 14:22:13 +0100
Subject: [PATCH] feat: add support for post-stages

---
 usr/lib/akshara/utils/gen_rootfs.py | 47 ++++++++++++++---------------
 usr/lib/akshara/utils/helpers.py    |  9 ++++++
 2 files changed, 31 insertions(+), 25 deletions(-)

diff --git a/usr/lib/akshara/utils/gen_rootfs.py b/usr/lib/akshara/utils/gen_rootfs.py
index 7384ee3..1f74b57 100644
--- a/usr/lib/akshara/utils/gen_rootfs.py
+++ b/usr/lib/akshara/utils/gen_rootfs.py
@@ -22,7 +22,7 @@ def gen_rootfs(system_config: dict, rootfs_path: str) -> RootFS:
         subprocess.run(
             ["bash", "-s"],
             text=True,
-            input=system_config["distro-config"]["before-stages"],
+            input=system_config["distro-config"]["initialise"],
             cwd=rootfs_path,
             env=os.environ.copy() | system_config["env"],
         ).returncode
@@ -33,33 +33,30 @@ def gen_rootfs(system_config: dict, rootfs_path: str) -> RootFS:
 
     modules = {}
 
-    if isinstance(system_config.get("modules"), list):
-        for module in system_config["modules"]:
-            modules[module["name"]] = module["run"]
+    for module in system_config["modules"]:
+        modules[module["name"]] = module["run"]
 
-    if isinstance(system_config.get("stages"), list):
-        for stage in system_config["stages"]:
-            if stage.get("module") not in modules.keys():
-                output.error(f"{stage.get('module')} not found within module list.")
-                exit(1)
-
-            inputs = stage["inputs"] if isinstance(stage.get("inputs"), list) else []
-            run_script_rootfs(rootfs, modules[stage["module"]], inputs)
-
-    if isinstance(system_config["distro-config"].get("after-stages"), str):
-        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")
+    for stage in system_config["stages"] + system_config["post-stages"]:
+        if stage.get("module") not in modules.keys():
+            output.error(f"{stage.get('module')} not found within module list.")
             exit(1)
 
+        inputs = stage["inputs"] if isinstance(stage.get("inputs"), list) else []
+        run_script_rootfs(rootfs, modules[stage["module"]], inputs)
+
+    if (
+        subprocess.run(
+            ["bash", "-s"],
+            text=True,
+            input=system_config["distro-config"]["finalise"],
+            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)
         pass
diff --git a/usr/lib/akshara/utils/helpers.py b/usr/lib/akshara/utils/helpers.py
index 779ceca..4f12652 100644
--- a/usr/lib/akshara/utils/helpers.py
+++ b/usr/lib/akshara/utils/helpers.py
@@ -19,6 +19,9 @@ def resolve_config(system_config: dict) -> dict:
                 "stages": system_config["stages"]
                 if isinstance(system_config.get("stages"), list)
                 else [],
+                "post-stages": system_config["post-stages"]
+                if isinstance(system_config.get("post-stages"), list)
+                else [],
                 "override": system_config["override"]
                 if isinstance(system_config.get("override"), list)
                 else [],
@@ -63,6 +66,12 @@ def resolve_config(system_config: dict) -> dict:
             else []
         )
 
+        base_config["post-stages"] = (
+            system_config["post-stages"]
+            if isinstance(system_config.get("post-stages"), list)
+            else []
+        ) + base_config["post-stages"]
+
         base_config["override"] += (
             system_config["override"]
             if isinstance(system_config.get("override"), list)