added functions to be imported into the init source, improved docs

Signed-off-by: Zen <[email protected]>

Author: desultory

readme.md

@@ -469,7 +469,14 @@ This module is loaded in the imports section of the `base.yaml` file:
 "ugrd.fs.mounts" = [ "_process_mounts_multi" ]
 ```
 
-#### build_tasks
+#### Imports
+
+UGRD allows functions to be imported from modules using the `imports` dict.
+
+This is primarily used to run additional functions at build time, add init functions, and add library functions.
+
+
+##### build_tasks
 
 Build tasks are functions which will be executed after the directory structure has been generated using the specified `paths`.
 
@@ -480,7 +487,7 @@ The base module includes a build task for generating the fstab, which is activat
 "ugrd.fs.mounts" = [ "generate_fstab" ]
 ```
 
-#### Packing tasks
+##### pack
 
 Packing facts, such as CPIO generation can be defined in the `pack` import.
 
@@ -491,7 +498,11 @@ The `cpio` module imports the `make_cpio_list` packing function with:
 "ugrd.base.base" = [ "make_cpio_list" ]
 ```
 
-#### init hooks
+##### funcs
+
+Functions can be added to `imports.funcs` to force the output to be added to `init_funcs.sh`.
+
+##### init hooks
 
 By default, the specified init hooks are:
 * `init_pre` - Where the base initramfs environment is set up, such as creating a devtmpfs.
@@ -514,7 +525,7 @@ The `InitramfsGenerator.generate_init_main()` function (often called from `self`
 
 A general overview of the procedure used for generating the init is to write the chosen `shebang`, then every init hook. The `custom_init` import can be used for more advanced confugrations, such as running another script in `agetty`.
 
-#### custom_init
+##### custom_init
 
 To change how everything but `init_pre` and `init_file` are handled at runtime, `custom_init` can be used.
 

ugrd/crypto/cryptsetup.py

@@ -1,5 +1,5 @@
 __author__ = 'desultory'
-__version__ = '1.0.1'
+__version__ = '1.1.0'
 
 _module_name = 'ugrd.crypto.cryptsetup'
 
@@ -97,10 +97,15 @@ def get_crypt_sources(self) -> list[str]:
             parameters['_host_device_path'] = _get_device_path_from_token(self, token)
         # Add a blkid command to get the source device in the initramfs, only match if the device has a partuuid
         out.append(f"export SOURCE_TOKEN_{name}='{token[0]}={token[1]}'")
-        out.append(f'export CRYPTSETUP_SOURCE_{name}=$(blkid --match-token "$SOURCE_TOKEN_{name}" --match-tag PARTUUID --output device)')
-
-        check_command = f'if [ -z "$CRYPTSETUP_SOURCE_{name}" ]; then echo "Unable to resolve device source for {name}"; bash; else echo "Resolved device source: $CRYPTSETUP_SOURCE_{name}"; fi\n'
-        out += [f"echo 'Attempting to get device path for {name}'", check_command]
+        source_cmd = f'export CRYPTSETUP_SOURCE_{name}=$(blkid --match-token "$SOURCE_TOKEN_{name}" --match-tag PARTUUID --output device)'
+
+        check_command = [f'if [ -z "$CRYPTSETUP_SOURCE_{name}" ]; then',
+                         f'    echo "Unable to resolve device source for {name}"',
+                         '    bash',
+                         'else',
+                         f'    echo "Resolved device source: $CRYPTSETUP_SOURCE_{name}"',
+                         'fi']
+        out += [f"echo 'Attempting to get device path for {name}'", source_cmd, *check_command]
 
     return out
 

ugrd/fs/btrfs.py

@@ -1,4 +1,4 @@
-__version__ = '0.1.0'
+__version__ = '0.2.0'
 __author__ = 'desultory'
 
 
@@ -23,23 +23,18 @@ def btrfs_scan(self) -> str:
     return "btrfs device scan"
 
 
-def _get_devices(self) -> str:
+def get_btrfs_devices(self) -> str:
     """
     Returns a bash function which uses blkid to get all btrfs devices
     """
-    return """
-    function get_btrfs_devices() {{
-        blkid -t TYPE=btrfs -o device
-    }}
-    """
+    return "blkid -t TYPE=btrfs -o device"
 
 
 def select_subvol(self) -> str:
     """
     selects a subvolume
     """
-    out = [_get_devices(self)]
-
+    out = []
     if root_subvol := self.config_dict.get("root_subvol"):
         out.append(f"btrfs subvolume set-default {root_subvol}")
 

ugrd/fs/btrfs.toml

@@ -3,6 +3,9 @@ _kmod_depend = [ "btrfs" ]
 
 subvol_selector = false
 
+[imports.functions]
+"ugrd.fs.btrfs" = [ "get_btrfs_devices" ]
+
 [imports.config_processing]
 "ugrd.fs.btrfs" = [ "_process_root_subvol" ]
 

ugrd/initramfs_dict.py

@@ -1,6 +1,6 @@
 
 __author__ = "desultory"
-__version__ = "0.9.0"
+__version__ = "0.10.0"
 
 from tomllib import load, TOMLDecodeError
 from pathlib import Path
@@ -123,6 +123,11 @@ def _process_imports(self, import_type: str, import_value: dict) -> None:
                     self['custom_processing'][function.__name__] = function
                     self.logger.debug("Registered config processing function: %s" % function.__name__)
 
+            if import_type == 'funcs':
+                for function in function_list:
+                    if function.__name__ in self['imports']['funcs']:
+                        raise ValueError("Function '%s' already registered" % function.__name__)
+
     @handle_plural
     def _process_mod_depends(self, module: str) -> None:
         """

ugrd/initramfs_generator.py

@@ -20,8 +20,8 @@ def __init__(self, config='/etc/ugrd/config.toml', *args, **kwargs):
         self.build_pre = [self.generate_structure]
         self.config_dict = InitramfsConfigDict(logger=self.logger)
 
-        # This will be filled with function definitions for imported functions
-        self._init_funcs = {}
+        # Used for functions that are added to the bash source file
+        self.included_functions = {}
 
         # init_pre and init_final are run as part of generate_initramfs_main
         self.init_types = ['init_debug', 'init_early', 'init_main', 'init_late', 'init_premount', 'init_mount', 'init_cleanup']
@@ -84,9 +84,12 @@ def build_structure(self) -> None:
         self._run_hook('build_pre', return_output=False)
         self._run_hook('build_tasks', return_output=False)
 
-    def _run_func(self, function, external=False, return_output=True) -> list[str]:
+    def _run_func(self, function, external=False, return_output=True, force_include=False) -> list[str]:
         """
-        Runs a function, returning the output in a list
+        Runs a function.
+        External ones must have self passed as the first argument.
+        Returns the output if return_output is set.
+        If force_include is set, forces the function to be included in the bash source file.
         """
         self.logger.debug("Running function: %s" % function.__name__)
         if external:
@@ -99,17 +102,27 @@ def _run_func(self, function, external=False, return_output=True) -> list[str]:
             return
 
         if function_output:
+            if isinstance(function_output, list) and len(function_output) == 1:
+                self.logger.debug("[%s] Function returned list with one element: %s" % (function.__name__, function_output[0]))
+                function_output = function_output[0]
+
+            if function.__name__ in self.included_functions:
+                raise ValueError("Function '%s' has already been included in the bash source file" % function.__name__)
+
             if isinstance(function_output, str):
                 self.logger.debug("[%s] Function returned string: %s" % (function.__name__, function_output))
+                if force_include:
+                    self.logger.debug("[%s] Function output is being forced into bash source file" % function.__name__)
+                    self.included_functions[function.__name__] = function_output
                 return function_output
             else:
                 self.logger.debug("[%s] Function returned output: %s" % (function.__name__, pretty_print(function_output)))
-                self._init_funcs[function.__name__] = function_output
+                self.included_functions[function.__name__] = function_output
                 return function.__name__
         else:
             self.logger.debug("[%s] Function returned no output" % function.__name__)
 
-    def _run_funcs(self, functions: list[str], external=False, return_output=True) -> list[str]:
+    def _run_funcs(self, functions: list[str], *args, return_output=True, **kwargs) -> list[str]:
         """
         Runs a list of functions
         Returns the output if return_output is set
@@ -119,25 +132,25 @@ def _run_funcs(self, functions: list[str], external=False, return_output=True) -
         for function in functions:
             # Only append if returning the output
             if return_output:
-                if function_output := self._run_func(function, external=external, return_output=return_output):
+                if function_output := self._run_func(function, *args, **kwargs):
                     out.append(function_output)
             else:
-                self._run_func(function, external=external, return_output=return_output)
+                self._run_func(function, *args, **kwargs)
         return out
 
-    def _run_hook(self, hook: str, return_output=True) -> list[str]:
+    def _run_hook(self, hook: str, *args, return_output=True, **kwargs) -> list[str]:
         """
         Runs a hook for imported functions
         """
         out = []
         self.logger.info("Running hook: %s" % hook)
         if hasattr(self, hook):
             self.logger.debug("Running internal functions for hook: %s" % hook)
-            out += self._run_funcs(getattr(self, hook), return_output=return_output)
+            out += self._run_funcs(getattr(self, hook), return_output=return_output, *args, **kwargs)
 
         if external_functions := self.config_dict['imports'].get(hook):
             self.logger.debug("Running external functions for hook: %s" % hook)
-            function_output = self._run_funcs(external_functions, external=True, return_output=return_output)
+            function_output = self._run_funcs(external_functions, external=True, return_output=return_output, *args, **kwargs)
             out += function_output
 
         if return_output:
@@ -154,18 +167,19 @@ def _run_init_hook(self, level: str) -> list[str]:
 
     def generate_init_funcs(self) -> None:
         """
-        Generates the init functions file based on self._init_funcs
+        Generates the init functions file based on self.included_functions
         """
-        if not self._init_funcs:
-            self.logger.warning("No init functions to generate")
-            return
-
         out = [f"# Generated by UGRD v{__version__}"]
 
-        for func_name, func_content in self._init_funcs.items():
+        for func_name, func_content in self.included_functions.items():
             out.append("\n\n" + func_name + "() {")
-            for line in func_content:
-                out.append(f"    {line}")
+            if isinstance(func_content, str):
+                out.append(f"    {func_content}")
+            elif isinstance(func_content, list):
+                for line in func_content:
+                    out.append(f"    {line}")
+            else:
+                raise TypeError("Function content is not a string or list: %s" % func_content)
             out.append("}")
 
         return out
@@ -194,6 +208,9 @@ def generate_init(self) -> None:
         init = [self.config_dict['shebang']]
         init += ["# Generated by UGRD v%s" % __version__]
 
+        # Run all included functions, so they get included
+        self._run_hook('functions', force_include=True)
+
         init.extend(self._run_init_hook('init_pre'))
 
         if self.config_dict['imports'].get('custom_init'):
@@ -205,7 +222,7 @@ def generate_init(self) -> None:
         init.extend(self._run_init_hook('init_final'))
         init += ["\n\n# END INIT"]
 
-        if self._init_funcs:
+        if self.included_functions:
             init_funcs = self.generate_init_funcs()
             self._write('init_funcs.sh', init_funcs, 0o755)
             init.insert(2, "source /init_funcs.sh")