From 2a08830ac0fa9d4d4a2491e7dd0c534cb9bf37cc Mon Sep 17 00:00:00 2001
From: Rudra Saraswat <rs2009@ruds.io>
Date: Thu, 2 Apr 2026 04:24:49 +0100
Subject: [PATCH] feat: add custom env support

---
 usr/lib/akshara/classes/rootfs.py   | 43 +++--------------------------
 usr/lib/akshara/utils/gen_rootfs.py |  6 ++--
 usr/lib/akshara/utils/helpers.py    |  9 ++++++
 3 files changed, 17 insertions(+), 41 deletions(-)

diff --git a/usr/lib/akshara/classes/rootfs.py b/usr/lib/akshara/classes/rootfs.py
index 1306a12..31c8238 100644
--- a/usr/lib/akshara/classes/rootfs.py
+++ b/usr/lib/akshara/classes/rootfs.py
@@ -10,15 +10,17 @@ class RootFS:
         rootfs_path: A string containing the path to the root filesystem
     """
 
-    def __init__(self, rootfs_path: str, distro_config: dict) -> None:
+    def __init__(self, rootfs_path: str, distro_config: dict, env: dict) -> None:
         """Initialises an instance based on a rootfs path and distro configuration.
 
         Args:
             rootfs: Path to root filesystem.
             distro_config: Dictionary containing distro configuration.
+            env: Dictionary containing environment for command execution.
         """
         self.rootfs_path = rootfs_path
         self.distro_config = distro_config
+        self.env = env
 
     def exists(self, path):
         """Checks if path exists within rootfs.
@@ -28,18 +30,6 @@ class RootFS:
         """
         return os.path.exists(os.path.join(self.rootfs_path, path))
 
-    def init(self) -> subprocess.CompletedProcess:
-        """Initialise rootfs."""
-
-        completedProcess = subprocess.run(
-            ["bash", "-s"],
-            text=True,
-            input=self.distro_config["before-stages"],
-            cwd=self.rootfs_path,
-        )
-
-        return completedProcess
-
     def exec(self, cmd, **kwargs) -> subprocess.CompletedProcess:
         """Runs command within rootfs.
 
@@ -52,6 +42,7 @@ class RootFS:
                 "systemd-nspawn",
                 "--quiet",
                 "--pipe",
+                *[f"--setenv={name}={val}" for name, val in self.env.items()],
                 "-D",
                 self.rootfs_path,
             ]
@@ -59,31 +50,5 @@ class RootFS:
             **kwargs,
         )
 
-    # def copy_kernels_to_boot(self) -> None:
-    #     """Copies any found kernels to /boot within rootfs."""
-    #     kernels = [
-    #         kernel
-    #         for kernel in os.listdir(f"{self.rootfs_path}/usr/lib/modules")
-    #         if self.exists(f"/usr/lib/modules/{kernel}/vmlinuz")
-    #     ]
-
-    #     if len(kernels) == 0:
-    #         return
-
-    #     for boot_file in os.listdir(f"{self.rootfs_path}/boot"):
-    #         if not os.path.isdir(boot_file):
-    #             self.exec(["rm", "-f", f"/boot/{boot_file}"])
-
-    #     for kernel in kernels:
-    #         self.exec(
-    #             ["cp", f"/usr/lib/modules/{kernel}/vmlinuz", f"/boot/vmlinuz-{kernel}"],
-    #             stdout=subprocess.DEVNULL,
-    #             stderr=subprocess.DEVNULL,
-    #         )
-
-    # def gen_initramfs(self) -> None:
-    #     """Generates initramfs within rootfs."""
-    #     self.exec(["dracut", "--force", "--regenerate-all"])
-
     def __repr__(self) -> str:
         return self.rootfs_path
diff --git a/usr/lib/akshara/utils/gen_rootfs.py b/usr/lib/akshara/utils/gen_rootfs.py
index eee1b1d..c68990e 100644
--- a/usr/lib/akshara/utils/gen_rootfs.py
+++ b/usr/lib/akshara/utils/gen_rootfs.py
@@ -13,16 +13,17 @@ def run_script_rootfs(rootfs: RootFS, input: str, args: list):
     rootfs.exec(["bash", "-s", *args], text=True, input=input)
 
 
-def gen_rootfs(system_config: dict, rootfs_path: str) -> RootFS:
+def gen_rootfs(system_config: dict, rootfs_path: str, env: dict = {}) -> RootFS:
     """Generates a rootfs for a given system configuration."""
 
-    rootfs = RootFS(rootfs_path, system_config["distro-config"])
+    rootfs = RootFS(rootfs_path, system_config["distro-config"], env)
 
     subprocess.run(
         ["bash", "-s"],
         text=True,
         input=system_config["distro-config"]["before-stages"],
         cwd=rootfs_path,
+        env=os.environ.copy() | env,
     )
 
     modules = {}
@@ -46,6 +47,7 @@ def gen_rootfs(system_config: dict, rootfs_path: str) -> RootFS:
             text=True,
             input=system_config["distro-config"]["after-stages"],
             cwd=str(rootfs_path),
+            env=os.environ.copy() | env,
         )
 
     with open(os.path.join(rootfs_path, "usr/system.json"), "w") as system_json_file:
diff --git a/usr/lib/akshara/utils/helpers.py b/usr/lib/akshara/utils/helpers.py
index e782eb7..874cc33 100644
--- a/usr/lib/akshara/utils/helpers.py
+++ b/usr/lib/akshara/utils/helpers.py
@@ -22,6 +22,9 @@ def resolve_config(system_config: dict) -> dict:
                 "override": system_config["override"]
                 if isinstance(system_config.get("override"), list)
                 else [],
+                "env": system_config["env"]
+                if isinstance(system_config.get("env"), dict)
+                else {},
                 "distro-config": system_config["distro-config"],
                 "auto-update": system_config["auto-update"]
                 if isinstance(system_config.get("auto-update"), dict)
@@ -63,6 +66,12 @@ def resolve_config(system_config: dict) -> dict:
             else []
         )
 
+        base_config["env"] = (
+            base_config["env"] | system_config["env"]
+            if isinstance(system_config.get("env"), dict)
+            else base_config["env"]
+        )
+
         base_config["auto-update"] = (
             system_config["auto-update"]
             if isinstance(system_config.get("auto-update"), dict)