From c1a006bb9ee915c08075e9adeab245f02452ee43 Mon Sep 17 00:00:00 2001
From: Martin Packman <martin.packman@canonical.com>
Date: Mon, 12 Mar 2012 13:56:40 +0000
Subject: Restrict warning to userdata without MIME wrapping only while still
 ignoring most types

---
 cloudinit/UserDataHandler.py     |  4 ++--
 cloudinit/__init__.py            | 14 ++++++-----
 tests/unittests/test_userdata.py | 50 ++++++++++++++++++++++++++++++++++++----
 3 files changed, 56 insertions(+), 12 deletions(-)

diff --git a/cloudinit/UserDataHandler.py b/cloudinit/UserDataHandler.py
index 98729056..ec914480 100644
--- a/cloudinit/UserDataHandler.py
+++ b/cloudinit/UserDataHandler.py
@@ -180,7 +180,7 @@ def process_includes(msg, appendmsg=None):
 
         payload = part.get_payload(decode=True)
 
-        if ctype_orig == "text/plain":
+        if ctype_orig in ("text/plain", "text/x-not-multipart"):
             ctype = type_from_startswith(payload)
 
         if ctype is None:
@@ -213,7 +213,7 @@ def message_from_string(data, headers=None):
             else:
                 msg[key] = val
     else:
-        mtype = headers.get("Content-Type", "text/plain")
+        mtype = headers.get("Content-Type", "text/x-not-multipart")
         maintype, subtype = mtype.split("/", 1)
         msg = MIMEBase(maintype, subtype, *headers)
         msg.set_payload(data)
diff --git a/cloudinit/__init__.py b/cloudinit/__init__.py
index 98a48fec..1c3f861c 100644
--- a/cloudinit/__init__.py
+++ b/cloudinit/__init__.py
@@ -606,12 +606,14 @@ def partwalker_callback(pdata, ctype, filename, payload):
         partwalker_handle_handler(pdata, ctype, filename, payload)
         return
     if ctype not in pdata['handlers']:
-        start = payload.split("\n", 1)[0][:24] # Use first line or 24 bytes
-        if start < payload:
-            details = "starting '%s...'" % start.encode("string-escape")
-        else:
-            details = repr(payload)
-        log.warning("Unhandled userdata part of type %s %s", ctype, details)
+        if ctype == "text/x-not-multipart":
+            # Extract the first line or 24 bytes for displaying in the log
+            start = payload.split("\n", 1)[0][:24]
+            if start < payload:
+                details = "starting '%s...'" % start.encode("string-escape")
+            else:
+                details = repr(payload)
+            log.warning("Unhandled non-multipart userdata %s", details)
         return
     handler_handle_part(pdata['handlers'][ctype], pdata['data'],
         ctype, filename, payload, pdata['frequency'])
diff --git a/tests/unittests/test_userdata.py b/tests/unittests/test_userdata.py
index 949f36dc..37ad9b13 100644
--- a/tests/unittests/test_userdata.py
+++ b/tests/unittests/test_userdata.py
@@ -3,6 +3,8 @@
 import logging
 import StringIO
 
+from email.mime.base import MIMEBase
+
 from mocker import MockerTestCase
 
 import cloudinit
@@ -42,7 +44,28 @@ class TestConsumeUserData(MockerTestCase):
         self._log = logging.getLogger(cloudinit.logger_name)
         self._log.addHandler(self._log_handler)
 
-    def test_script(self):
+    def test_unhandled_type_warning(self):
+        """Raw text without magic is ignored but shows warning"""
+        self.mocker.replay()
+        ci = cloudinit.CloudInit()
+        ci.datasource = FakeDataSource("arbitrary text\n")
+        ci.consume_userdata()
+        self.assertEqual(
+            "Unhandled non-multipart userdata starting 'arbitrary text...'\n",
+            self.log_file.getvalue())
+
+    def test_mime_text_plain(self):
+        """Mime message of type text/plain is ignored without warning"""
+        self.mocker.replay()
+        ci = cloudinit.CloudInit()
+        message = MIMEBase("text", "plain")
+        message.set_payload("Just text")
+        ci.datasource = FakeDataSource(message.as_string())
+        ci.consume_userdata()
+        self.assertEqual("", self.log_file.getvalue())
+
+    def test_shellscript(self):
+        """Raw text starting #!/bin/sh is treated as script"""
         script = "#!/bin/sh\necho hello\n"
         outpath = cloudinit.get_ipath_cur("scripts") + "/part-001"
         self.mock_write(outpath, script, 0700)
@@ -52,9 +75,28 @@ class TestConsumeUserData(MockerTestCase):
         ci.consume_userdata()
         self.assertEqual("", self.log_file.getvalue())
 
-    def test_unhandled_type_warning(self):
+    def test_mime_text_x_shellscript(self):
+        """Mime message of type text/x-shellscript is treated as script"""
+        script = "#!/bin/sh\necho hello\n"
+        outpath = cloudinit.get_ipath_cur("scripts") + "/part-001"
+        self.mock_write(outpath, script, 0700)
         self.mocker.replay()
         ci = cloudinit.CloudInit()
-        ci.datasource = FakeDataSource("arbitrary text\n")
+        message = MIMEBase("text", "x-shellscript")
+        message.set_payload(script)
+        ci.datasource = FakeDataSource(message.as_string())
         ci.consume_userdata()
-        self.assertIn("Unhandled userdata part", self.log_file.getvalue())
+        self.assertEqual("", self.log_file.getvalue())
+
+    def test_mime_text_plain_shell(self):
+        """Mime type text/plain starting #!/bin/sh is treated as script"""
+        script = "#!/bin/sh\necho hello\n"
+        outpath = cloudinit.get_ipath_cur("scripts") + "/part-001"
+        self.mock_write(outpath, script, 0700)
+        self.mocker.replay()
+        ci = cloudinit.CloudInit()
+        message = MIMEBase("text", "plain")
+        message.set_payload(script)
+        ci.datasource = FakeDataSource(message.as_string())
+        ci.consume_userdata()
+        self.assertEqual("", self.log_file.getvalue())
-- 
cgit v1.2.3