From ca9257cf232ec74a3908464ddcbf060f2e7869dd Mon Sep 17 00:00:00 2001
From: Alex Youngs <alexyoungs@hatchbed.com>
Date: Wed, 30 Oct 2024 18:07:55 -0500
Subject: [PATCH] Added bashrc_extesions extension

---
 README.md                                     |  1 +
 setup.py                                      |  1 +
 src/rocker/bashrc_extension.py                | 79 +++++++++++++++++++
 ...shrc_extensions_user_snippet.Dockerfile.em |  6 ++
 4 files changed, 87 insertions(+)
 create mode 100644 src/rocker/bashrc_extension.py
 create mode 100644 src/rocker/templates/bashrc_extensions_user_snippet.Dockerfile.em

diff --git a/README.md b/README.md
index 38f14f2..edc847d 100644
--- a/README.md
+++ b/README.md
@@ -37,6 +37,7 @@ You can get full details on the extensions from the main `rocker --help` command
 - home -- Mount the user's home directory into the container
 - pulse -- Mount pulse audio into the container
 - ssh -- Pass through ssh access to the container.
+- bashrc_extensions -- Pass a list of files/URLs of bash scripts that will extend the container's native .bashrc
 
 As well as access to many of the docker arguments as well such as `device`, `env`, `volume`, `name`, `network`, `ipc`, and `privileged`.
 
diff --git a/setup.py b/setup.py
index 2cc1d9d..9224ba0 100644
--- a/setup.py
+++ b/setup.py
@@ -45,6 +45,7 @@
             'detect_docker_image_os = rocker.cli:detect_image_os',
         ],
         'rocker.extensions': [
+            'bashrc_extensions = rocker.bashrc_extension:BashrcExtensions',
             'cuda = rocker.nvidia_extension:Cuda',
             'devices = rocker.extensions:Devices',
             'dev_helpers = rocker.extensions:DevHelpers',
diff --git a/src/rocker/bashrc_extension.py b/src/rocker/bashrc_extension.py
new file mode 100644
index 0000000..1cbc202
--- /dev/null
+++ b/src/rocker/bashrc_extension.py
@@ -0,0 +1,79 @@
+import em
+import pkgutil
+import os
+import urllib3
+from rocker.core import get_user_name, ExtensionError
+from rocker.extensions import RockerExtension
+
+
+class BashrcExtensions(RockerExtension):
+
+    name = 'bashrc_extensions'
+
+    @classmethod
+    def get_name(cls):
+        return cls.name
+
+    def __init__(self):
+        self._env_subs = None
+        self.name = BashrcExtensions.get_name()
+
+    def precondition_environment(self, cli_args):
+        pass
+
+    def validate_environment(self, cli_args):
+        pass
+
+    def get_preamble(self, cli_args):
+        return ''
+    
+    def get_filename(self, bashrc_extension_file):
+        if os.path.isfile(bashrc_extension_file):
+            return os.path.join(self.name, os.path.basename(bashrc_extension_file))
+        elif bashrc_extension_file.startswith(('http://', 'https://')):
+            filename = os.path.basename(urllib3.util.url.parse_url(bashrc_extension_file).path)
+            if not filename:
+                raise ExtensionError('Bashrc extension file does not appear to have a filename: {}'.format(bashrc_extension_file))
+            return os.path.join(self.name, filename)
+        else:
+            raise ExtensionError('Bashrc extension files is not a file or URL: {}'.format(bashrc_extension_file))
+
+    def get_files(self, cli_args):
+        files = {}
+        for bashrc_extension in cli_args[self.name]:
+            if os.path.isfile(bashrc_extension):
+                with open(bashrc_extension, 'r') as f:
+                    files[self.get_filename(bashrc_extension)] = f.read()
+            elif bashrc_extension.startswith(('http://', 'https://')):
+                try:
+                    response = urllib3.PoolManager().request('GET', bashrc_extension)
+                    if response.status != 200:
+                        raise ExtensionError(f'Failed to fetch bashrc extension from URL {bashrc_extension}, status code: {response.status}')
+                    files[self.get_filename(bashrc_extension)] = response.data.decode('utf-8')
+                except urllib3.exceptions.HTTPError as e:
+                    raise ExtensionError(f'Failed to fetch bashrc extension from URL {bashrc_extension}: {str(e)}')
+
+        return files
+
+    @staticmethod
+    def get_home_dir(cli_args):
+        if cli_args["user"]:
+            return os.path.join(os.path.sep, "home", get_user_name())
+        else:
+            return os.path.join(os.path.sep, "root")
+
+    def get_user_snippet(self, cli_args):
+        args = {}
+        args['bashrc_extension_files'] = {self.get_filename(bashrc_extension): os.path.basename(self.get_filename(bashrc_extension)) for bashrc_extension in cli_args[self.name]}
+        args['home_dir'] = self.get_home_dir(cli_args)
+
+        snippet = pkgutil.get_data(
+            'rocker', 'templates/{}_user_snippet.Dockerfile.em'.format(self.name)).decode('utf-8')
+
+        return em.expand(snippet, args)
+
+    @staticmethod
+    def register_arguments(parser, defaults={}):
+        parser.add_argument('--bashrc-extensions',
+                            nargs='+',
+                            help="Sources custom bashrc extensions from the container's default bashrc. An extension can be a local file or a URL.")
diff --git a/src/rocker/templates/bashrc_extensions_user_snippet.Dockerfile.em b/src/rocker/templates/bashrc_extensions_user_snippet.Dockerfile.em
new file mode 100644
index 0000000..26eca08
--- /dev/null
+++ b/src/rocker/templates/bashrc_extensions_user_snippet.Dockerfile.em
@@ -0,0 +1,6 @@
+WORKDIR @home_dir
+RUN echo "\n# Source custom bashrc extensions" >> @home_dir/.bashrc
+@[for path, filename in bashrc_extension_files.items()]
+COPY @path @filename
+RUN echo ". ~/@filename" >> @home_dir/.bashrc
+@[end for]