diff options
| -rwxr-xr-x | ec2-run-user-data.py | 80 | ||||
| -rw-r--r-- | tests.py | 101 |
2 files changed, 173 insertions, 8 deletions
diff --git a/ec2-run-user-data.py b/ec2-run-user-data.py index 5af04249..67ecf219 100755 --- a/ec2-run-user-data.py +++ b/ec2-run-user-data.py @@ -22,6 +22,7 @@ import email import os import subprocess import tempfile +from xml.dom.minidom import parse, parseString import ec2init @@ -34,14 +35,6 @@ def register_handler(mimetype, func): content_type_handlers[mimetype] = func return func -def main(): - ec2 = ec2init.EC2Init() - - user_data = ec2.get_user_data() - - msg = email.message_from_string(user_data) - handle_part(msg) - def handle_part(part): if part.is_multipart(): for p in part.get_payload(): @@ -57,6 +50,14 @@ def handle_unknown_payload(payload): # Try to detect magic if payload.startswith('#!'): content_type_handlers['text/x-shellscript'](payload) + return + if payload.startswith('<appliance>'): + content_type_handlers['text/x-appliance-config'](payload) + +@handler('text/x-appliance-config') +def handle_appliance_config(payload): + app = ApplianceConfig(payload) + app.handle() @handler('text/x-ebs-mount-description') def handle_ebs_mount_description(payload): @@ -87,5 +88,68 @@ def handle_shell_script(payload): os.unlink(path) +class ApplianceConfig(object): + def __init__(self, data): + self.data = data + + def handle(self): + self.dom = parseString(self.data) + + if self.dom.childNodes[0].tagName == 'appliance': + root = self.dom.childNodes[0] + else: + return + + for node in root.childNodes: + if node.tagName == 'package': + pkg = None + for subnode in node.childNodes: + if subnode.nodeType == root.TEXT_NODE: + pkg = subnode.nodeValue + if not pkg: + # Something's fishy. We should have been passed the name of + # a package. + return + if node.getAttribute('action') == 'remove': + remove_package(pkg) + else: + install_package(pkg) + elif node.tagName == 'script': + script = '' + for subnode in node.childNodes: + # If someone went through the trouble of wrapping it in CDATA, + # it's probably the script we want to run.. + if subnode.nodeType == root.CDATA_SECTION_NODE: + script = subnode.nodeValue + # ..however, fall back to whatever TEXT_NODE stuff is between + # the <script> tags. + if subnode.nodeType == root.TEXT_NODE and not script: + script = subnode.nodeValue + if not script: + # An empty script? + continue + content_type_handlers['text/x-shellscript'](script) + +def main(): + ec2 = ec2init.EC2Init() + + user_data = ec2.get_user_data() + msg = parse_user_data(user_data) + handle_part(msg) + +def parse_user_data(user_data): + return email.message_from_string(user_data) + +def install_remove_package(pkg, action): + apt_get = subprocess.Popen(['apt-get', action, pkg], stdout=subprocess.PIPE) + logger_process = subprocess.Popen(['logger', '-t', 'user-data'], stdin=apt_get.stdout) + logger_process.communicate() + +def install_package(pkg): + return install_remove_package(pkg, 'install') + +def remove_package(pkg): + return install_remove_package(pkg, 'remove') + if __name__ == '__main__': main() diff --git a/tests.py b/tests.py new file mode 100644 index 00000000..33c45ec7 --- /dev/null +++ b/tests.py @@ -0,0 +1,101 @@ +#!/usr/bin/python +# +# Unit tests for EC2-init +# Copyright (C) 2008-2009 Canonical Ltd. +# +# Author: Soren Hansen <soren@canonical.com> +# +# This program is free software: you can redistribute it and/or modify +# it under the terms of the GNU General Public License version 3, as +# published by the Free Software Foundation. +# +# This program is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU General Public License for more details. +# +# You should have received a copy of the GNU General Public License +# along with this program. If not, see <http://www.gnu.org/licenses/>. +# + +import unittest + +class RunUserDataApplianceConfigScript(unittest.TestCase): + def setUp(self): + self.ec2_run_user_data = __import__('ec2-run-user-data') + self.fake_handle_shell_script_counter = 0 + self.expected_scripts = [] + # Override install_remove_package + self.ec2_run_user_data.content_type_handlers['text/x-shellscript'] = self.fake_handle_shell_script + + def fake_handle_shell_script(self, txt): + self.fake_handle_shell_script_counter += 1 + self.assertEqual(self.expected_scripts.pop(0), txt) + + def handle_xml(self, xml): + msg = self.ec2_run_user_data.parse_user_data(xml) + self.ec2_run_user_data.handle_part(msg) + + def testApplianceConfigPackageScriptSingle(self): + script = '''#!/bin/sh +echo hey''' + xml = '<appliance><script>%s</script></appliance>' % script + self.expected_scripts += [script] + self.handle_xml(xml) + self.assertEqual(self.fake_handle_shell_script_counter, 1) + + def testApplianceConfigPackageScriptMultiple(self): + script1 = '''#!/bin/sh +echo hey''' + script2 = '''#!/usr/bin/python +print "hey"''' + xml = '<appliance><script>%s</script><script>%s</script></appliance>' % (script1, script2) + self.expected_scripts += [script1, script2] + self.handle_xml(xml) + self.assertEqual(self.fake_handle_shell_script_counter, 2) + + def testApplianceConfigPackageScriptCDATA(self): + script = '''#!/bin/sh +echo hey''' + xml = '<appliance><script><![CDATA[%s]]></script></appliance>' % (script, ) + self.expected_scripts += [script] + self.handle_xml(xml) + self.assertEqual(self.fake_handle_shell_script_counter, 1) + +class RunUserDataApplianceConfigPackageHandling(unittest.TestCase): + def setUp(self): + self.fake_install_remove_package_counter = 0 + + self.ec2_run_user_data = __import__('ec2-run-user-data') + + # Override install_remove_package + self.ec2_run_user_data.install_remove_package = self.fake_install_remove_package + + def fake_install_remove_package(self, package, action): + self.fake_install_remove_package_counter += 1 + mapping = { 'foobarplus': 'install', + 'foobarminus': 'remove' } + self.assert_(package in mapping) + self.assertEqual(action, mapping[package]) + + def handle_xml(self, xml): + msg = self.ec2_run_user_data.parse_user_data(xml) + self.ec2_run_user_data.handle_part(msg) + + def testApplianceConfigPackageInstall(self): + xml = '<appliance><package>foobarplus</package></appliance>' + self.handle_xml(xml) + self.assertEqual(self.fake_install_remove_package_counter, 1) + + def testApplianceConfigPackageRemove(self): + xml = '<appliance><package action="remove">foobarminus</package></appliance>' + self.handle_xml(xml) + self.assertEqual(self.fake_install_remove_package_counter, 1) + + def testApplianceConfigPackageInstallAndRemove(self): + xml = '<appliance><package>foobarplus</package><package action="remove">foobarminus</package></appliance>' + self.handle_xml(xml) + self.assertEqual(self.fake_install_remove_package_counter, 2) + +if __name__ == "__main__": + unittest.main() |