Merge pull request #1768 from nimbis/ssm-backend-get-command-invocation-upstream

Ssm backend get command invocation upstream

moto/ssm/models.py

@@ -5,10 +5,12 @@ from collections import defaultdict
 from moto.core import BaseBackend, BaseModel
 from moto.core.exceptions import RESTError
 from moto.ec2 import ec2_backends
+from moto.cloudformation import cloudformation_backends
 
 import datetime
 import time
 import uuid
+import itertools
 
 
 class Parameter(BaseModel):
@@ -67,7 +69,7 @@ class Command(BaseModel):
             instance_ids=None, max_concurrency='', max_errors='',
             notification_config=None, output_s3_bucket_name='',
             output_s3_key_prefix='', output_s3_region='', parameters=None,
-            service_role_arn='', targets=None):
+            service_role_arn='', targets=None, backend_region='us-east-1'):
 
         if instance_ids is None:
             instance_ids = []
@@ -88,9 +90,9 @@ class Command(BaseModel):
         self.status = 'Success'
         self.status_details = 'Details placeholder'
 
-        now = datetime.datetime.now()
+        self.requested_date_time = datetime.datetime.now()
-        self.requested_date_time = now.isoformat()
+        self.requested_date_time_iso = self.requested_date_time.isoformat()
-        expires_after = now + datetime.timedelta(0, timeout_seconds)
+        expires_after = self.requested_date_time + datetime.timedelta(0, timeout_seconds)
         self.expires_after = expires_after.isoformat()
 
         self.comment = comment
@@ -105,6 +107,32 @@ class Command(BaseModel):
         self.parameters = parameters
         self.service_role_arn = service_role_arn
         self.targets = targets
+        self.backend_region = backend_region
+
+        # Get instance ids from a cloud formation stack target.
+        stack_instance_ids = [self.get_instance_ids_by_stack_ids(target['Values']) for
+            target in self.targets if
+            target['Key'] == 'tag:aws:cloudformation:stack-name']
+
+        self.instance_ids += list(itertools.chain.from_iterable(stack_instance_ids))
+
+        # Create invocations with a single run command plugin.
+        self.invocations = []
+        for instance_id in self.instance_ids:
+            self.invocations.append(
+                self.invocation_response(instance_id, "aws:runShellScript"))
+
+    def get_instance_ids_by_stack_ids(self, stack_ids):
+        instance_ids = []
+        cloudformation_backend = cloudformation_backends[self.backend_region]
+        for stack_id in stack_ids:
+            stack_resources = cloudformation_backend.list_stack_resources(stack_id)
+            instance_resources = [
+                instance.id for instance in stack_resources
+                if instance.type == "AWS::EC2::Instance"]
+            instance_ids.extend(instance_resources)
+
+        return instance_ids
 
     def response_object(self):
         r = {
@@ -122,7 +150,7 @@ class Command(BaseModel):
             'OutputS3BucketName': self.output_s3_bucket_name,
             'OutputS3KeyPrefix': self.output_s3_key_prefix,
             'Parameters': self.parameters,
-            'RequestedDateTime': self.requested_date_time,
+            'RequestedDateTime': self.requested_date_time_iso,
             'ServiceRole': self.service_role_arn,
             'Status': self.status,
             'StatusDetails': self.status_details,
@@ -132,6 +160,50 @@ class Command(BaseModel):
 
         return r
 
+    def invocation_response(self, instance_id, plugin_name):
+        # Calculate elapsed time from requested time and now. Use a hardcoded
+        # elapsed time since there is no easy way to convert a timedelta to
+        # an ISO 8601 duration string.
+        elapsed_time_iso = "PT5M"
+        elapsed_time_delta = datetime.timedelta(minutes=5)
+        end_time = self.requested_date_time + elapsed_time_delta
+
+        r = {
+            'CommandId': self.command_id,
+            'InstanceId': instance_id,
+            'Comment': self.comment,
+            'DocumentName': self.document_name,
+            'PluginName': plugin_name,
+            'ResponseCode': 0,
+            'ExecutionStartDateTime': self.requested_date_time_iso,
+            'ExecutionElapsedTime': elapsed_time_iso,
+            'ExecutionEndDateTime': end_time.isoformat(),
+            'Status': 'Success',
+            'StatusDetails': 'Success',
+            'StandardOutputContent': '',
+            'StandardOutputUrl': '',
+            'StandardErrorContent': '',
+        }
+
+        return r
+
+    def get_invocation(self, instance_id, plugin_name):
+        invocation = next(
+            (invocation for invocation in self.invocations
+                if invocation['InstanceId'] == instance_id), None)
+
+        if invocation is None:
+            raise RESTError(
+                'InvocationDoesNotExist',
+                'An error occurred (InvocationDoesNotExist) when calling the GetCommandInvocation operation')
+
+        if plugin_name is not None and invocation['PluginName'] != plugin_name:
+                raise RESTError(
+                    'InvocationDoesNotExist',
+                    'An error occurred (InvocationDoesNotExist) when calling the GetCommandInvocation operation')
+
+        return invocation
+
 
 class SimpleSystemManagerBackend(BaseBackend):
 
@@ -140,6 +212,11 @@ class SimpleSystemManagerBackend(BaseBackend):
         self._resource_tags = defaultdict(lambda: defaultdict(dict))
         self._commands = []
 
+        # figure out what region we're in
+        for region, backend in ssm_backends.items():
+            if backend == self:
+                self._region = region
+
     def delete_parameter(self, name):
         try:
             del self._parameters[name]
@@ -260,7 +337,8 @@ class SimpleSystemManagerBackend(BaseBackend):
             output_s3_region=kwargs.get('OutputS3Region', ''),
             parameters=kwargs.get('Parameters', {}),
             service_role_arn=kwargs.get('ServiceRoleArn', ''),
-            targets=kwargs.get('Targets', []))
+            targets=kwargs.get('Targets', []),
+            backend_region=self._region)
 
         self._commands.append(command)
         return {
@@ -298,6 +376,18 @@ class SimpleSystemManagerBackend(BaseBackend):
             command for command in self._commands
             if instance_id in command.instance_ids]
 
+    def get_command_invocation(self, **kwargs):
+        """
+        https://docs.aws.amazon.com/systems-manager/latest/APIReference/API_GetCommandInvocation.html
+        """
+
+        command_id = kwargs.get('CommandId')
+        instance_id = kwargs.get('InstanceId')
+        plugin_name = kwargs.get('PluginName', None)
+
+        command = self.get_command_by_id(command_id)
+        return command.get_invocation(instance_id, plugin_name)
+
 
 ssm_backends = {}
 for region, ec2_backend in ec2_backends.items():

moto/ssm/responses.py

@@ -210,3 +210,8 @@ class SimpleSystemManagerResponse(BaseResponse):
         return json.dumps(
             self.ssm_backend.list_commands(**self.request_params)
         )
+
+    def get_command_invocation(self):
+        return json.dumps(
+            self.ssm_backend.get_command_invocation(**self.request_params)
+        )

tests/test_ssm/test_ssm_boto3.py

@@ -5,11 +5,12 @@ import botocore.exceptions
 import sure   # noqa
 import datetime
 import uuid
+import json
 
 from botocore.exceptions import ClientError
 from nose.tools import assert_raises
 
-from moto import mock_ssm
+from moto import mock_ssm, mock_cloudformation
 
 
 @mock_ssm
@@ -668,3 +669,118 @@ def test_list_commands():
     with assert_raises(ClientError):
         response = client.list_commands(
             CommandId=str(uuid.uuid4()))
+
+@mock_ssm
+def test_get_command_invocation():
+    client = boto3.client('ssm', region_name='us-east-1')
+
+    ssm_document = 'AWS-RunShellScript'
+    params = {'commands': ['#!/bin/bash\necho \'hello world\'']}
+
+    response = client.send_command(
+        InstanceIds=['i-123456', 'i-234567', 'i-345678'],
+        DocumentName=ssm_document,
+        Parameters=params,
+        OutputS3Region='us-east-2',
+        OutputS3BucketName='the-bucket',
+        OutputS3KeyPrefix='pref')
+
+    cmd = response['Command']
+    cmd_id = cmd['CommandId']
+
+    instance_id = 'i-345678'
+    invocation_response = client.get_command_invocation(
+        CommandId=cmd_id,
+        InstanceId=instance_id,
+        PluginName='aws:runShellScript')
+
+    invocation_response['CommandId'].should.equal(cmd_id)
+    invocation_response['InstanceId'].should.equal(instance_id)
+
+    # test the error case for an invalid instance id
+    with assert_raises(ClientError):
+        invocation_response = client.get_command_invocation(
+            CommandId=cmd_id,
+            InstanceId='i-FAKE')
+
+    # test the error case for an invalid plugin name
+    with assert_raises(ClientError):
+        invocation_response = client.get_command_invocation(
+            CommandId=cmd_id,
+            InstanceId=instance_id,
+            PluginName='FAKE')
+
+@mock_ssm
+@mock_cloudformation
+def test_get_command_invocations_from_stack():
+    stack_template = {
+        "AWSTemplateFormatVersion": "2010-09-09",
+        "Description": "Test Stack",
+        "Resources": {
+            "EC2Instance1": {
+                "Type": "AWS::EC2::Instance",
+                "Properties": {
+                    "ImageId": "ami-test-image-id",
+                    "KeyName": "test",
+                    "InstanceType": "t2.micro",
+                    "Tags": [
+                        {
+                            "Key": "Test Description",
+                            "Value": "Test tag"
+                        },
+                        {
+                            "Key": "Test Name",
+                            "Value": "Name tag for tests"
+                        }
+                    ]
+                }
+            }
+        },
+        "Outputs": {
+            "test": {
+                "Description": "Test Output",
+                "Value": "Test output value",
+                "Export": {
+                    "Name": "Test value to export"
+                }
+            },
+            "PublicIP": {
+                "Value": "Test public ip"
+            }
+        }
+    }
+
+    cloudformation_client = boto3.client(
+        'cloudformation',
+        region_name='us-east-1')
+
+    stack_template_str = json.dumps(stack_template)
+
+    response = cloudformation_client.create_stack(
+        StackName='test_stack',
+        TemplateBody=stack_template_str,
+        Capabilities=('CAPABILITY_IAM', ))
+
+    client = boto3.client('ssm', region_name='us-east-1')
+
+    ssm_document = 'AWS-RunShellScript'
+    params = {'commands': ['#!/bin/bash\necho \'hello world\'']}
+
+    response = client.send_command(
+        Targets=[{
+            'Key': 'tag:aws:cloudformation:stack-name',
+            'Values': ('test_stack', )}],
+        DocumentName=ssm_document,
+        Parameters=params,
+        OutputS3Region='us-east-2',
+        OutputS3BucketName='the-bucket',
+        OutputS3KeyPrefix='pref')
+
+    cmd = response['Command']
+    cmd_id = cmd['CommandId']
+    instance_ids = cmd['InstanceIds']
+
+    invocation_response = client.get_command_invocation(
+        CommandId=cmd_id,
+        InstanceId=instance_ids[0],
+        PluginName='aws:runShellScript')