Repository: JonathonReinhart/scuba
Pull request #58 (2016-05-14 21:23:37)

Implement scubainit hooks
See #54


Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 0c7f34d..0822e6e 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -3,6 +3,9 @@ All notable changes to this project will be documented in this file.
 This project adheres to [Semantic Versioning](http://semver.org/).
 
 ## [Unreleased]
+### Added
+- Add support for scubainit hooks
+
 ### Changed
 - `scubainit` re-implemented as a C program, which does the following:
    - Creates the scubauser user/group
diff --git a/README.md b/README.md
index 01fd09a..3753f8a 100644
--- a/README.md
+++ b/README.md
@@ -56,19 +56,39 @@ Required nodes:
 Optional nodes:
 
 - `aliases` - A dictionary of bash-like aliases
+- `hooks` - "Hooks" that run as part of `scubainit` before running the user command
 
 An example `.scuba.yml` file might look like this:
 
 ```yaml
 image: gcc:5.1
+
 aliases:
   build: make -j4
+
+
+# These hooks run during scubainit
+hooks:
+
+  # The "root" hook runs just before we switch users
+  root:
+    # This uses the complex form, with a 'script' subkey,
+    # which is a list of command strings making up the script.
+    script:
+      - 'echo "HOOK: This runs before we switch users"'
+      - id
+
+  # The "user" hook runs just before we execute the user command.
+  # This uses the simple form, which is just a single command string.
+  user: 'echo "HOOK: After switching users, uid=$(id -u) gid=$(id -g)"'
 ```
 
 This tells SCUBA:
 - Use the `gcc:5.1` Docker image
 - `build` is an alias for `make -j4`.
+
 In this example, `scuba build foo` would execute `make -j4 foo` in a `gcc:5.1` container.
+Prior to running `make`, the hooks would execute, echoing messages about the user ID.
 
 ### Extended syntax
 In addition to normal YAML syntax, an additional constructor, `!from_yaml`, is available for `.scuba.yml` which allows a key to be retrieved from an external YAML file. Is has the following syntax:
diff --git a/example/run_all.sh b/example/run_all.sh
index 34c5d2a..c734038 100755
--- a/example/run_all.sh
+++ b/example/run_all.sh
@@ -1,9 +1,9 @@
 #!/bin/sh
 cd $(dirname $0)
-export PATH="$(realpath ../src):$PATH"
 
 ./01_basic/run_example.sh
 ./02_with_alias/run_example.sh
 ./03_search_up/one/two/three/four/run_example.sh
 ./external_yaml_simple/run_example.sh
 ./external_yaml_nested/run_example.sh
+./scubainit_hooks/run_example.sh
diff --git a/example/scubainit_hooks/.scuba.yml b/example/scubainit_hooks/.scuba.yml
new file mode 100644
index 0000000..7e48ec5
--- /dev/null
+++ b/example/scubainit_hooks/.scuba.yml
@@ -0,0 +1,16 @@
+image: busybox:latest
+
+# These hooks run during scubainit
+hooks:
+
+  # The "root" hook runs just before we switch users
+  root:
+    # This uses the complex form, with a 'script' subkey,
+    # which is a list of command strings making up the script.
+    script:
+      - 'echo "HOOK: This runs before we switch users"'
+      - id
+
+  # The "user" hook runs just before we execute the user command.
+  # This uses the simple form, which is just a single command string.
+  user: 'echo "HOOK: After switching users, uid=$(id -u) gid=$(id -g)"'
diff --git a/example/scubainit_hooks/run_example.sh b/example/scubainit_hooks/run_example.sh
new file mode 100755
index 0000000..2161788
--- /dev/null
+++ b/example/scubainit_hooks/run_example.sh
@@ -0,0 +1,5 @@
+#!/bin/bash
+
+cd $(dirname $0)
+
+scuba echo "Finally, the output from the user command"
diff --git a/scuba/__main__.py b/scuba/__main__.py
index 94a9409..51843cf 100644
--- a/scuba/__main__.py
+++ b/scuba/__main__.py
@@ -33,6 +33,28 @@ def get_umask():
     os.umask(val)
     return val
 
+def generate_hook_script(config, opts, name):
+    script = config.hooks.get(name)
+    if not script:
+        return
+
+    def writeln(f, line):
+        f.write(line + '\n')
+
+    # Generate the hook script, mount it into the container, and tell scubainit
+    with NamedTemporaryFile(mode='wt', prefix='scuba', delete=False) as f:
+        filecleanup.register(f.name)
+
+        cpath = '/.scuba/hooks/{0}.sh'.format(name)
+        opts.append(make_vol_opt(f.name, cpath))
+        opts.append('--env=SCUBAINIT_HOOK_{0}={1}'.format(name.upper(), cpath))
+
+        writeln(f, '#!/bin/sh')
+        writeln(f, '# Auto-generated from .scuba.yml')
+        writeln(f, 'set -e')
+        for cmd in script:
+            writeln(f, cmd)
+
 
 def get_native_opts(config, scuba_args, usercmd):
     opts = [
@@ -71,6 +93,10 @@ def get_native_opts(config, scuba_args, usercmd):
             sys.exit(128)
         verbose_msg('{0} Cmd: "{1}"'.format(config.image, usercmd))
 
+    # Hooks
+    for name in ('root', 'user', ):
+        generate_hook_script(config, opts, name)
+
     return opts, usercmd
 
 
diff --git a/scuba/config.py b/scuba/config.py
index dd7e77c..250b92e 100644
--- a/scuba/config.py
+++ b/scuba/config.py
@@ -1,6 +1,10 @@
 import os
 import yaml
 import shlex
+try:
+    basestring
+except NameError:
+    basestring = str    # Python 3
 
 from .constants import *
 
@@ -102,7 +106,7 @@ def find_config():
 class ScubaConfig(object):
     def __init__(self, **data):
         required_nodes = ('image',)
-        optional_nodes = ('aliases',)
+        optional_nodes = ('aliases','hooks',)
 
         # Check for missing required nodes
         missing = [n for n in required_nodes if not n in data]
@@ -122,6 +126,59 @@ class ScubaConfig(object):
         for alias, cmdstr in data.get('aliases', {}).items():
             self._aliases[alias] = shlex_split(cmdstr)
 
+        self._load_hooks(data)
+
+
+    def _process_script(self, node, name):
+        '''Process a script-type node
+
+        This can handle yaml of either a simple form:
+
+            node: this is my script
+
+        Or a more complex form (which allows for other sub-nodes):
+
+            node:
+              script:
+                - this is my script
+                - it has multiple parts
+
+        Other forms are disallowed:
+
+            node:
+              - this
+              - is
+              - forbidden
+        '''
+        if isinstance(node, basestring):
+            # The script is just the text itself
+            return [node]
+
+
+        if isinstance(node, dict):
+            # There must be a "script" key, which must be a list of strings
+            script = node.get('script')
+            if not script:
+                raise ConfigError("{0}: must have a 'script' subkey".format(name))
+
+            if not isinstance(script, list):
+                raise ConfigError("{0}.script: must be a list".format(name))
+
+            return script
+
+        raise ConfigError("{0}: must be string or dict".format(name))
+
+
+    def _load_hooks(self, data):
+        self._hooks = {}
+
+        for name in ('user', 'root',):
+            node = data.get('hooks', {}).get(name)
+            if node:
+                hook = self._process_script(node, name)
+                self._hooks[name] = hook
+
+
     @property
     def image(self):
         return self._image
@@ -130,6 +187,11 @@ class ScubaConfig(object):
     def aliases(self):
         return self._aliases
 
+    @property
+    def hooks(self):
+        return self._hooks
+
+
     def process_command(self, command):
         if command:
             rep = self.aliases.get(command[0])
diff --git a/scubainit/scubainit.c b/scubainit/scubainit.c
index e9dce3d..5929128 100644
--- a/scubainit/scubainit.c
+++ b/scubainit/scubainit.c
@@ -35,7 +35,8 @@
 #define SCUBAINIT_UID       "SCUBAINIT_UID"
 #define SCUBAINIT_GID       "SCUBAINIT_GID"
 #define SCUBAINIT_UMASK     "SCUBAINIT_UMASK"
-#define SCUBAINIT_HOOK      "SCUBAINIT_HOOK"
+#define SCUBAINIT_HOOK_USER "SCUBAINIT_HOOK_USER"
+#define SCUBAINIT_HOOK_ROOT "SCUBAINIT_HOOK_ROOT"
 #define SCUBAINIT_VERBOSE   "SCUBAINIT_VERBOSE"
 
 static bool m_verbose = false;
@@ -43,6 +44,9 @@ static int m_uid = -1;
 static int m_gid = -1;
 static int m_umask = -1;
 
+static const char *m_user_hook;
+static const char *m_root_hook;
+
 
 /* Returns true if the scuba user should be used */
 static bool
@@ -361,6 +365,59 @@ make_homedir(const char *path, unsigned int uid, unsigned int gid)
     return 0;
 }
 
+static int
+make_executable(const char *path)
+{
+    int ret = -1;
+    struct stat st;
+    mode_t mode;
+
+    if (stat(path, &st) != 0)
+        goto out;
+
+    mode = st.st_mode;
+
+    /* Copy R bits to X */
+    mode |= (mode & 0444) >> 2;
+
+    if (chmod(path, mode) != 0)
+        goto out;
+
+    ret = 0;
+
+out:
+    return ret;
+}
+
+static int
+call_hook(const char *hook_path)
+{
+    int rc;
+
+    /* Nothing to do? */
+    if (hook_path == NULL)
+        return 0;
+
+    if (make_executable(hook_path) != 0) {
+        errmsg("Failed to make executable %s: %m\n", hook_path);
+        exit(99);
+    }
+
+    verbose("About to execute %s\n", hook_path);
+    rc = system(hook_path);
+
+    if (rc < 0) {
+        errmsg("Failed to execute %s: %m\n", hook_path);
+        exit(99);
+    }
+    if (rc > 0) {
+        errmsg("%s exited with status %d\n", hook_path, rc);
+        exit(rc);
+    }
+
+    return 0;
+}
+
 static int
 str2uint(const char *val, unsigned int *result)
 {
@@ -416,6 +473,24 @@ getenv_uint_opt_unset(const char *name, int *result)
     return 0;
 }
 
+/**
+ * Gets an environment variable string, and unsets it
+ */
+static char *
+getenv_str_unset(const char *name)
+{
+    char *var;
+
+    if (!(var = getenv(name)))
+        return NULL;
+
+    /* Duplicate the string then unset the var */
+    var = strdup(var);
+    unsetenv(name);
+
+    return var;
+}
+
 
 static int
 process_envvars(void)
@@ -466,6 +541,11 @@ process_envvars(void)
         m_verbose = true;
     }
 
+    /* Hook scripts */
+    m_user_hook = getenv_str_unset(SCUBAINIT_HOOK_USER);
+    m_root_hook = getenv_str_unset(SCUBAINIT_HOOK_ROOT);
+
+
     /* Clear out other env. vars */
     unsetenv("PWD");
     unsetenv("OLDPWD");
@@ -507,7 +587,8 @@ main(int argc, char **argv)
             exit(99);
     }
 
-    /* FUTURE: Call pre-su hook */
+    /* Call pre-su hook */
+    call_hook(m_root_hook);
 
     /* Handle the scuba user */
     if (use_scuba_user()) {
@@ -520,7 +601,10 @@ main(int argc, char **argv)
         umask(m_umask);
     }
 
-    /* FUTURE: Call post-su hook */
+    /* Call post-su hook, only if we switch users */
+    if (use_scuba_user()) {
+        call_hook(m_user_hook);
+    }
 
 
     /* Prepare for execution of user command */


Test patch:
diff --git a/tests/test_config.py b/tests/test_config.py
index 502545d..fb14e0d 100644
--- a/tests/test_config.py
+++ b/tests/test_config.py
@@ -225,3 +225,67 @@ class TestConfig(TestCase):
                 )
         result = cfg.process_command(['apple', 'arg1', 'arg2'])
         assert_equal(result, ['banana', 'cherry', 'pie is good', 'arg1', 'arg2'])
+
+    ############################################################################
+    # Hooks
+
+    def test_hooks_mixed(self):
+        '''hooks of mixed forms are valid'''
+        with open('.scuba.yml', 'w') as f:
+            f.write('''
+                image: na
+                hooks:
+                  root:
+                    script:
+                      - echo "This runs before we switch users"
+                      - id
+                  user: id
+                ''')
+
+        config = scuba.config.load_config('.scuba.yml')
+
+        assert_seq_equal(
+            config.hooks.get('root'),
+            ['echo "This runs before we switch users"', 'id'])
+
+        assert_seq_equal(
+            config.hooks.get('user'),
+            ['id'])
+
+    def test_hooks_invalid_list(self):
+        '''hooks with list not under "script" key are invalid'''
+        with open('.scuba.yml', 'w') as f:
+            f.write('''
+                image: na
+                hooks:
+                  user:
+                    - this list should be under
+                    - a 'script'
+                ''')
+
+        assert_raises(scuba.config.ConfigError, scuba.config.load_config, '.scuba.yml')
+
+    # TODO: Any reason this shouldn't be valid?
+    def test_hooks_invalid_script_type(self):
+        '''hooks with string "script" are invalid'''
+        with open('.scuba.yml', 'w') as f:
+            f.write('''
+                image: na
+                hooks:
+                  user:
+                    script: this should be in a list under script
+                ''')
+
+        assert_raises(scuba.config.ConfigError, scuba.config.load_config, '.scuba.yml')
+
+    def test_hooks_missing_script(self):
+        '''hooks with dict, but missing "script" are invalid'''
+        with open('.scuba.yml', 'w') as f:
+            f.write('''
+                image: na
+                hooks:
+                  user:
+                    not_script: missing "script" key
+                ''')
+
+        assert_raises(scuba.config.ConfigError, scuba.config.load_config, '.scuba.yml')
diff --git a/tests/test_main.py b/tests/test_main.py
index 66f4bb5..7207a98 100644
--- a/tests/test_main.py
+++ b/tests/test_main.py
@@ -342,3 +342,32 @@ class TestMain(TestCase):
             out, _ = self.run_scuba(args)
 
         assert_str_equalish(out, data)
+
+
+    ############################################################################
+    # Hooks
+
+    def _test_one_hook(self, hookname, exp_uid, exp_gid):
+        with open('.scuba.yml', 'w') as f:
+            f.write('image: {0}\n'.format(DOCKER_IMAGE))
+            f.write('hooks:\n')
+            f.write('  {0}: echo $(id -u) $(id -g)\n'.format(hookname))
+
+        args = ['/bin/sh', '-c', 'echo success']
+        out, _ = self.run_scuba(args)
+
+        out = out.splitlines()
+
+        uid, gid = map(int, out[0].split())
+        assert_equal(exp_uid, uid)
+        assert_equal(exp_gid, gid)
+
+        assert_str_equalish(out[1], 'success')
+
+    def test_user_hook(self):
+        '''Verify user hook executes as user'''
+        self._test_one_hook('user', os.getuid(), os.getgid())
+
+    def test_root_hook(self):
+        '''Verify root hook executes as root'''
+        self._test_one_hook('root', 0, 0)