Merge pull request #107 from flomotlik/feature/main-account

Feature/main account

formica/__init__.py

@@ -1,7 +1,7 @@
 import logging
 import sys
 
-__version__ = '0.8.11'
+__version__ = '0.8.12'
 
 CHANGE_SET_FORMAT = "{stack}-change-set"
 

formica/cli.py

@@ -30,7 +30,8 @@
     'capabilities': list,
     'vars': dict,
     'administration_role_arn': str,
-    'execution_role_name': str
+    'execution_role_name': str,
+    'main_account': bool
 }
 
 
@@ -184,6 +185,7 @@ def stack_set_parser(parser):
     add_aws_arguments(create_parser)
     add_stack_set_argument(create_parser)
     add_stack_parameters_argument(create_parser)
+    add_stack_set_main_account_parameter(create_parser)
     add_stack_tags_argument(create_parser)
     add_capabilities_argument(create_parser)
     add_config_file_argument(create_parser)
@@ -195,6 +197,7 @@ def stack_set_parser(parser):
     add_aws_arguments(update_parser)
     add_stack_set_argument(update_parser)
     add_stack_parameters_argument(update_parser)
+    add_stack_set_main_account_parameter(update_parser)
     add_stack_tags_argument(update_parser)
     add_capabilities_argument(update_parser)
     add_config_file_argument(update_parser)
@@ -289,6 +292,10 @@ def add_stack_set_instance_retain_argument(parser):
     parser.add_argument('--retain', help='Retain stacks', action='store_true', default=False)
 
 
+def add_stack_set_main_account_parameter(parser):
+    parser.add_argument('--main-account', help='Set MainAccount Parameter', action='store_true', default=False)
+
+
 def add_config_file_argument(parser):
     parser.add_argument('--config-file', '-c',
                         type=argparse.FileType('r'),

formica/stack_set.py

@@ -1,6 +1,7 @@
 from formica.aws import AWS
 import logging
 import sys
+import json
 
 logger = logging.getLogger(__name__)
 
@@ -67,7 +68,14 @@ def remove_stack_set_instances(args):
 def __manage_stack_set(args, create):
     from .loader import Loader
     client = AWS.current_session().client('cloudformation')
-    params = parameters(parameters=args.parameters,
+    params = args.parameters or {}
+    main_account = args.main_account
+    if main_account:
+        sts = AWS.current_session().client('sts')
+        identity = sts.get_caller_identity()
+        params['MainAccount'] = identity['Account']
+
+    params = parameters(parameters=params,
                         tags=args.tags,
                         capabilities=args.capabilities,
                         accounts=vars(args).get('accounts'),
@@ -77,17 +85,24 @@ def __manage_stack_set(args, create):
 
     loader = Loader(variables=args.vars)
     loader.load()
+    template = loader.template()
+    if main_account:
+        template = loader.template_dictionary()
+        template['Parameters'] = template.get('Parameters') or {}
+        template['Parameters']['MainAccount'] = {'Type': 'String'}
+        template = json.dumps(template)
+
     if create:
         result = client.create_stack_set(
             StackSetName=args.stack_set,
-            TemplateBody=loader.template(),
+            TemplateBody=template,
             ** params
         )
         logger.info('StackSet {} created'.format(args.stack_set))
     else:
         result = client.update_stack_set(
             StackSetName=args.stack_set,
-            TemplateBody=loader.template(),
+            TemplateBody=template,
             ** params
         )
         logger.info('StackSet {} updated in Operation {}'.format(args.stack_set, result['OperationId']))

tests/unit/test_stack_set.py

@@ -1,4 +1,6 @@
 import pytest
+import json
+from uuid import uuid4
 
 from formica import cli
 from tests.unit.constants import STACK, CLOUDFORMATION_PARAMETERS, CLOUDFORMATION_TAGS, TEMPLATE
@@ -64,6 +66,49 @@ def test_create_stack_set_without_arguments(client, logger, loader):
     )
 
 
+def test_create_stack_set_with_main_account(session, client, logger, mocker):
+    mock = mocker.patch('formica.loader.Loader')
+    mock.return_value.template_dictionary.return_value = {}
+    accountid = str(uuid4())
+    client.get_caller_identity.return_value = {'Account': accountid}
+    cli.main([
+        'stack-set',
+        'create',
+        '--stack-set', STACK,
+        '--main-account'
+    ])
+
+    session.return_value.client.assert_called_with('sts')
+
+    client.create_stack_set.assert_called_with(
+        StackSetName=STACK,
+        TemplateBody=json.dumps({'Parameters': {'MainAccount': {'Type': 'String'}}}),
+        Parameters=[{'ParameterKey': 'MainAccount', 'ParameterValue': accountid, 'UsePreviousValue': False}]
+    )
+
+
+def test_create_stack_set_with_main_account_and_existing_parameters(session, client, logger, mocker):
+    mock = mocker.patch('formica.loader.Loader')
+    mock.return_value.template_dictionary.return_value = {'Parameters': {'SomeParam': {'Type': 'String'}}}
+    accountid = str(uuid4())
+    client.get_caller_identity.return_value = {'Account': accountid}
+    cli.main([
+        'stack-set',
+        'create',
+        '--stack-set', STACK,
+        '--main-account',
+        '--parameters', 'A=B'
+
+    ])
+
+    client.create_stack_set.assert_called_with(
+        StackSetName=STACK,
+        TemplateBody=json.dumps({'Parameters': {'SomeParam': {'Type': 'String'}, 'MainAccount': {'Type': 'String'}}}),
+        Parameters=[{'ParameterKey': 'A', 'ParameterValue': 'B', 'UsePreviousValue': False}, {
+            'ParameterKey': 'MainAccount', 'ParameterValue': accountid, 'UsePreviousValue': False}]
+    )
+
+
 def test_update_stack_set(client, loader):
     client.update_stack_set.return_value = {'OperationId': '12345'}
 
@@ -93,6 +138,28 @@ def test_update_stack_set(client, loader):
     )
 
 
+def test_update_stack_set_with_main_account(session, client, logger, mocker):
+    client.update_stack_set.return_value = {'OperationId': '12345'}
+    mock = mocker.patch('formica.loader.Loader')
+    mock.return_value.template_dictionary.return_value = {}
+    accountid = str(uuid4())
+    client.get_caller_identity.return_value = {'Account': accountid}
+    cli.main([
+        'stack-set',
+        'update',
+        '--stack-set', STACK,
+        '--main-account'
+    ])
+
+    session.return_value.client.assert_called_with('sts')
+
+    client.update_stack_set.assert_called_with(
+        StackSetName=STACK,
+        TemplateBody=json.dumps({'Parameters': {'MainAccount': {'Type': 'String'}}}),
+        Parameters=[{'ParameterKey': 'MainAccount', 'ParameterValue': accountid, 'UsePreviousValue': False}]
+    )
+
+
 def test_add_stack_set_instances(client, loader):
     cli.main([
         'stack-set',