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.

4 files changed, 105 insertions, 22 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),
         ]