Prevent senza patch from destroying user data

senza/error_handling.py

@@ -15,7 +15,8 @@
 from raven import Client
 
 from .configuration import configuration
-from .exceptions import InvalidDefinition, PiuNotFound, SecurityGroupNotFound
+from .exceptions import (InvalidDefinition, InvalidUserDataType,
+                         PiuNotFound, SecurityGroupNotFound)
 from .manaus.exceptions import (ELBNotFound, HostedZoneNotFound, InvalidState,
                                 RecordNotFound)
 from .manaus.utils import extract_client_error_code
@@ -134,7 +135,7 @@ def __call__(self, *args, **kwargs):
                 "{}\nYou can install piu with the following command:"
                 "\nsudo pip3 install --upgrade stups-piu".format(error))
         except (ELBNotFound, HostedZoneNotFound, RecordNotFound,
-                InvalidDefinition, InvalidState) as error:
+                InvalidDefinition, InvalidState, InvalidUserDataType) as error:
             die_fatal_error(error)
         except SecurityGroupNotFound as error:
             message = ("{}\nRun `senza init` to (re-)create "

senza/exceptions.py

@@ -53,3 +53,29 @@ def __init__(self, security_group: str):
 
     def __str__(self):
         return 'Security Group "{}" does not exist.'.format(self.security_group)
+
+
+class InvalidUserDataType(SenzaException):
+    """
+    Exception raised when the type of the new user data is different from the
+    old user data
+    """
+
+    def __init__(self, old_type: type, new_type: type):
+        self.old_type = old_type
+        self.new_type = new_type
+
+    def __str__(self):
+        return ('Current user data is a {} but provided user data '
+                'is a {}.').format(self.__human_readable_type(self.old_type),
+                                   self.__human_readable_type(self.new_type))
+
+    def __human_readable_type(self, t) -> str:
+        if t is str:
+            return "string"
+        elif t is dict:
+            return "map"
+        elif t is int:
+            return 'integer'
+        else:
+            return str(t)

senza/patch.py

@@ -4,6 +4,7 @@
 
 import yaml
 
+from .exceptions import InvalidUserDataType
 from .manaus.boto_proxy import BotoClientProxy
 
 LAUNCH_CONFIGURATION_PROPERTIES = set([
@@ -57,8 +58,14 @@ def patch_auto_scaling_group(group: dict, region: str, properties: dict):
             now = datetime.datetime.utcnow().strftime('%Y%m%dT%H%M%S')
             kwargs['LaunchConfigurationName'] = '{}-{}'.format(kwargs['LaunchConfigurationName'][:64], now)
             for key, val in properties.items():
-                if key == 'UserData' and isinstance(val, dict):
-                    kwargs[key] = patch_user_data(kwargs[key], val)
+
+                if key == 'UserData':
+                    current_user_data = yaml.safe_load(kwargs['UserData'])
+                    if isinstance(val, dict):
+                        kwargs[key] = patch_user_data(kwargs[key], val)
+                    elif isinstance(current_user_data, dict):
+                        raise InvalidUserDataType(type(current_user_data),
+                                                  type(val))
                 else:
                     kwargs[key] = val
             asg.create_launch_configuration(**kwargs)

tests/test_patch.py

@@ -1,7 +1,9 @@
 import codecs
 
 from unittest.mock import MagicMock
+import pytest
 
+from senza.exceptions import InvalidUserDataType
 from senza.patch import patch_auto_scaling_group
 
 def test_patch_auto_scaling_group(monkeypatch):
@@ -48,3 +50,29 @@ def create_lc(**kwargs):
     patch_auto_scaling_group(group, 'myregion', properties)
 
     assert new_lc['UserData'] == '#firstline\nsource: newsource\n'
+
+
+
+def test_patch_user_data_wrong_type(monkeypatch):
+
+    lc = {'ImageId': 'originalimage', 'LaunchConfigurationName': 'originallc',
+          'UserData': codecs.encode(b'#firstline\nsource: oldsource', 'base64').decode('utf-8')}
+    result = {'LaunchConfigurations': [lc]}
+
+    asg = MagicMock()
+    asg.describe_launch_configurations.return_value = result
+
+    new_lc = {}
+
+    def create_lc(**kwargs):
+        new_lc.update(kwargs)
+
+    asg.create_launch_configuration = create_lc
+    monkeypatch.setattr('boto3.client', lambda x, region: asg)
+    group = {'AutoScalingGroupName': 'myasg', 'LaunchConfigurationName': 'originallc'}
+    properties = {'UserData': "it's a string"}
+    with pytest.raises(InvalidUserDataType) as exc_info:
+        patch_auto_scaling_group(group, 'myregion', properties)
+
+    assert str(exc_info.value) == ('Current user data is a map but provided '
+                                   'user data is a string.')