Shell script handlers by freq (#1166)

Handlers for per-boot/per-instance/per-once multipart MIME

Add handlers for adding scripts to userdata that can be run at various
frequencies. Scripts of type x-shellscript-per-boot,
x-shellscript-per-instance, or x-shellscript-per-once can be added
to a multipart MIME userdata message as part of instance userdata.
These scripts will then be added to the appropriate per-boot,
per-instance, or per-once directory in /var/lib/cloud/scripts/
during processing of userdata.

7 files changed, 178 insertions, 23 deletions

diff --git a/cloudinit/cmd/devel/make_mime.py b/cloudinit/cmd/devel/make_mime.py
index a7493c74..c7671a93 100755
--- a/cloudinit/cmd/devel/make_mime.py
+++ b/cloudinit/cmd/devel/make_mime.py
@@ -20,6 +20,28 @@ EPILOG = (
 )
 
 
+def create_mime_message(files):
+    sub_messages = []
+    errors = []
+    for i, (fh, filename, format_type) in enumerate(files):
+        contents = fh.read()
+        sub_message = MIMEText(contents, format_type, sys.getdefaultencoding())
+        sub_message.add_header(
+            "Content-Disposition", 'attachment; filename="%s"' % (filename)
+        )
+        content_type = sub_message.get_content_type().lower()
+        if content_type not in get_content_types():
+            msg = (
+                "content type %r for attachment %s " "may be incorrect!"
+            ) % (content_type, i + 1)
+            errors.append(msg)
+        sub_messages.append(sub_message)
+    combined_message = MIMEMultipart()
+    for msg in sub_messages:
+        combined_message.attach(msg)
+    return (combined_message, errors)
+
+
 def file_content_type(text):
     """Return file content type by reading the first line of the input."""
     try:
@@ -97,29 +119,14 @@ def handle_args(name, args):
         print("\n".join(get_content_types(strip_prefix=True)))
         return 0
 
-    sub_messages = []
-    errors = []
-    for i, (fh, filename, format_type) in enumerate(args.files):
-        contents = fh.read()
-        sub_message = MIMEText(contents, format_type, sys.getdefaultencoding())
-        sub_message.add_header(
-            "Content-Disposition", 'attachment; filename="%s"' % (filename)
-        )
-        content_type = sub_message.get_content_type().lower()
-        if content_type not in get_content_types():
-            level = "WARNING" if args.force else "ERROR"
-            msg = (
-                level + ": content type %r for attachment %s may be incorrect!"
-            ) % (content_type, i + 1)
-            sys.stderr.write(msg + "\n")
-            errors.append(msg)
-        sub_messages.append(sub_message)
-    if len(errors) and not args.force:
+    combined_message, errors = create_mime_message(args.files)
+    if errors:
+        level = "WARNING" if args.force else "ERROR"
+        for error in errors:
+            sys.stderr.write(f"{level}: {error}\n")
         sys.stderr.write("Invalid content-types, override with --force\n")
-        return 1
-    combined_message = MIMEMultipart()
-    for msg in sub_messages:
-        combined_message.attach(msg)
+        if not args.force:
+            return 1
     print(combined_message)
     return 0
 
diff --git a/cloudinit/handlers/__init__.py b/cloudinit/handlers/__init__.py
index 62c1fd26..7d8a9208 100644
--- a/cloudinit/handlers/__init__.py
+++ b/cloudinit/handlers/__init__.py
@@ -50,6 +50,13 @@ INCLUSION_TYPES_MAP = {
     "#cloud-config-archive": "text/cloud-config-archive",
     "#cloud-config-jsonp": "text/cloud-config-jsonp",
     "## template: jinja": "text/jinja2",
+    # Note: for the next 3 entries, the prefix doesn't matter because these
+    # are for types that can only be used as part of a MIME message. However,
+    # including these entries supresses warnings during `cloudinit devel
+    # make-mime`, which otherwise would require `--force`.
+    "text/x-shellscript-per-boot": "text/x-shellscript-per-boot",
+    "text/x-shellscript-per-instance": "text/x-shellscript-per-instance",
+    "text/x-shellscript-per-once": "text/x-shellscript-per-once",
 }
 
 # Sorted longest first
diff --git a/cloudinit/handlers/shell_script_by_frequency.py b/cloudinit/handlers/shell_script_by_frequency.py
new file mode 100644
index 00000000..923cca57
--- /dev/null
+++ b/cloudinit/handlers/shell_script_by_frequency.py
@@ -0,0 +1,62 @@
+import os
+
+from cloudinit import log, util
+from cloudinit.handlers import Handler
+from cloudinit.settings import PER_ALWAYS, PER_INSTANCE, PER_ONCE
+
+LOG = log.getLogger(__name__)
+
+# cloudinit/settings.py defines PER_*** frequency constants. It makes sense to
+# use them here, instead of hardcodes, and map them to the 'per-***' frequency-
+# specific folders in /v/l/c/scripts. It might make sense to expose this at a
+# higher level or in a more general module -- eg maybe in cloudinit/settings.py
+# itself -- but for now it's here.
+path_map = {
+    PER_ALWAYS: "per-boot",
+    PER_INSTANCE: "per-instance",
+    PER_ONCE: "per-once",
+}
+
+
+def get_mime_type_by_frequency(freq):
+    mime_type = f"text/x-shellscript-{path_map[freq]}"
+    return mime_type
+
+
+def get_script_folder_by_frequency(freq, scripts_dir):
+    """Return the frequency-specific subfolder for a given frequency constant
+    and parent folder."""
+    freqPath = path_map[freq]
+    folder = os.path.join(scripts_dir, freqPath)
+    return folder
+
+
+def write_script_by_frequency(script_path, payload, frequency, scripts_dir):
+    """Given a filename, a payload, a frequency, and a scripts folder, write
+    the payload to the correct frequency-specific path"""
+    filename = os.path.basename(script_path)
+    filename = util.clean_filename(filename)
+    folder = get_script_folder_by_frequency(frequency, scripts_dir)
+    path = os.path.join(folder, filename)
+    payload = util.dos2unix(payload)
+    util.write_file(path, payload, 0o700)
+
+
+class ShellScriptByFreqPartHandler(Handler):
+    """Common base class for the frequency-specific script handlers."""
+
+    def __init__(self, script_frequency, paths, **_kwargs):
+        Handler.__init__(self, PER_ALWAYS)
+        self.prefixes = [get_mime_type_by_frequency(script_frequency)]
+        self.script_frequency = script_frequency
+        self.scripts_dir = paths.get_cpath("scripts")
+        if "script_path" in _kwargs:
+            self.scripts_dir = paths.get_cpath(_kwargs["script_path"])
+
+    def handle_part(self, data, ctype, script_path, payload, frequency):
+        if script_path is not None:
+            filename = os.path.basename(script_path)
+            filename = util.clean_filename(filename)
+            write_script_by_frequency(
+                script_path, payload, self.script_frequency, self.scripts_dir
+            )
diff --git a/cloudinit/stages.py b/cloudinit/stages.py
index b1a6bc49..3f17294b 100644
--- a/cloudinit/stages.py
+++ b/cloudinit/stages.py
@@ -21,12 +21,16 @@ from cloudinit.handlers.boot_hook import BootHookPartHandler
 from cloudinit.handlers.cloud_config import CloudConfigPartHandler
 from cloudinit.handlers.jinja_template import JinjaTemplatePartHandler
 from cloudinit.handlers.shell_script import ShellScriptPartHandler
+from cloudinit.handlers.shell_script_by_frequency import (
+    ShellScriptByFreqPartHandler,
+)
 from cloudinit.handlers.upstart_job import UpstartJobPartHandler
 from cloudinit.net import cmdline
 from cloudinit.reporting import events
 from cloudinit.settings import (
     CLOUD_CONFIG,
     FREQUENCIES,
+    PER_ALWAYS,
     PER_INSTANCE,
     PER_ONCE,
     RUN_CLOUD_CONFIG,
@@ -519,6 +523,9 @@ class Init(object):
         def_handlers = [
             cloudconfig_handler,
             shellscript_handler,
+            ShellScriptByFreqPartHandler(PER_ALWAYS, **opts),
+            ShellScriptByFreqPartHandler(PER_INSTANCE, **opts),
+            ShellScriptByFreqPartHandler(PER_ONCE, **opts),
             BootHookPartHandler(**opts),
             UpstartJobPartHandler(**opts),
         ]
diff --git a/tests/integration_tests/test_shell_script_by_frequency.py b/tests/integration_tests/test_shell_script_by_frequency.py
new file mode 100644
index 00000000..25157722
--- /dev/null
+++ b/tests/integration_tests/test_shell_script_by_frequency.py
@@ -0,0 +1,48 @@
+"""Integration tests for various handlers."""
+
+from io import StringIO
+
+import pytest
+
+from cloudinit.cmd.devel.make_mime import create_mime_message
+from tests.integration_tests.instances import IntegrationInstance
+
+PER_FREQ_TEMPLATE = """\
+#!/bin/bash
+touch /tmp/test_per_freq_{}
+"""
+
+PER_ALWAYS_FILE = StringIO(PER_FREQ_TEMPLATE.format("always"))
+PER_INSTANCE_FILE = StringIO(PER_FREQ_TEMPLATE.format("instance"))
+PER_ONCE_FILE = StringIO(PER_FREQ_TEMPLATE.format("once"))
+
+FILES = [
+    (PER_ALWAYS_FILE, "always.sh", "x-shellscript-per-boot"),
+    (PER_INSTANCE_FILE, "instance.sh", "x-shellscript-per-instance"),
+    (PER_ONCE_FILE, "once.sh", "x-shellscript-per-once"),
+]
+
+USER_DATA, errors = create_mime_message(FILES)
+
+
+@pytest.mark.ci
+@pytest.mark.user_data(USER_DATA)
+def test_per_freq(client: IntegrationInstance):
+    # Sanity test for scripts folder
+    cmd = "test -d /var/lib/cloud/scripts"
+    assert client.execute(cmd).ok
+    # Test per-boot
+    cmd = "test -f /var/lib/cloud/scripts/per-boot/always.sh"
+    assert client.execute(cmd).ok
+    cmd = "test -f /tmp/test_per_freq_always"
+    assert client.execute(cmd).ok
+    # Test per-instance
+    cmd = "test -f /var/lib/cloud/scripts/per-instance/instance.sh"
+    assert client.execute(cmd).ok
+    cmd = "test -f /tmp/test_per_freq_instance"
+    assert client.execute(cmd).ok
+    # Test per-once
+    cmd = "test -f /var/lib/cloud/scripts/per-once/once.sh"
+    assert client.execute(cmd).ok
+    cmd = "test -f /tmp/test_per_freq_once"
+    assert client.execute(cmd).ok
diff --git a/tests/unittests/test_builtin_handlers.py b/tests/unittests/test_builtin_handlers.py
index a057be2a..0dae924d 100644
--- a/tests/unittests/test_builtin_handlers.py
+++ b/tests/unittests/test_builtin_handlers.py
@@ -12,6 +12,7 @@ from textwrap import dedent
 import pytest
 
 from cloudinit import handlers, helpers, subp, util
+from cloudinit.cmd.devel import read_cfg_paths
 from cloudinit.handlers.cloud_config import CloudConfigPartHandler
 from cloudinit.handlers.jinja_template import (
     JinjaTemplatePartHandler,
@@ -19,8 +20,12 @@ from cloudinit.handlers.jinja_template import (
     render_jinja_payload,
 )
 from cloudinit.handlers.shell_script import ShellScriptPartHandler
+from cloudinit.handlers.shell_script_by_frequency import (
+    get_script_folder_by_frequency,
+    path_map,
+)
 from cloudinit.handlers.upstart_job import UpstartJobPartHandler
-from cloudinit.settings import PER_ALWAYS, PER_INSTANCE
+from cloudinit.settings import PER_ALWAYS, PER_INSTANCE, PER_ONCE
 from tests.unittests.helpers import (
     CiTestCase,
     FilesystemMockingTestCase,
@@ -473,4 +478,22 @@ class TestRenderJinjaPayload(CiTestCase):
         self.assertIn(expected_log, self.logs.getvalue())
 
 
+class TestShellScriptByFrequencyHandlers:
+    def do_test_frequency(self, frequency):
+        ci_paths = read_cfg_paths()
+        scripts_dir = ci_paths.get_cpath("scripts")
+        testFolder = os.path.join(scripts_dir, path_map[frequency])
+        folder = get_script_folder_by_frequency(frequency, scripts_dir)
+        assert testFolder == folder
+
+    def test_get_script_folder_per_boot(self):
+        self.do_test_frequency(PER_ALWAYS)
+
+    def test_get_script_folder_per_instance(self):
+        self.do_test_frequency(PER_INSTANCE)
+
+    def test_get_script_folder_per_once(self):
+        self.do_test_frequency(PER_ONCE)
+
+
 # vi: ts=4 expandtab
diff --git a/tools/.github-cla-signers b/tools/.github-cla-signers
index 6d957d9b..ac157a2f 100644
--- a/tools/.github-cla-signers
+++ b/tools/.github-cla-signers
@@ -8,6 +8,7 @@ andrewbogott
 andrewlukoshko
 antonyc
 aswinrajamannar
+beantaxi
 beezly
 bipinbachhao
 BirknerAlex