From 84960d13db8c7db1ca9c94b56cf3ced00919cc36 Mon Sep 17 00:00:00 2001
From: Shay Nehmad <shay.nehmad@guardicore.com>
Date: Wed, 10 Apr 2019 20:03:12 +0300
Subject: [PATCH] Refactored and tested the "get_instances" method.

---
 monkey/common/cloud/aws_service.py            | 31 ++++++----
 ..._filter_instance_data_from_aws_response.py | 56 +++++++++++++++++++
 2 files changed, 75 insertions(+), 12 deletions(-)
 create mode 100644 monkey/common/cloud/test_filter_instance_data_from_aws_response.py

diff --git a/monkey/common/cloud/aws_service.py b/monkey/common/cloud/aws_service.py
index 735c4e976..522056385 100644
--- a/monkey/common/cloud/aws_service.py
+++ b/monkey/common/cloud/aws_service.py
@@ -6,6 +6,21 @@ from common.cloud.aws_instance import AwsInstance
 
 __author__ = 'itay.mizeretz'
 
+INSTANCE_INFORMATION_LIST_KEY = 'InstanceInformationList'
+INSTANCE_ID_KEY = 'InstanceId'
+COMPUTER_NAME_KEY = 'ComputerName'
+PLATFORM_TYPE_KEY = 'PlatformType'
+IP_ADDRESS_KEY = 'IPAddress'
+
+
+def filter_instance_data_from_aws_response(response):
+    return [{
+        'instance_id': x[INSTANCE_ID_KEY],
+        'name': x[COMPUTER_NAME_KEY],
+        'os': x[PLATFORM_TYPE_KEY].lower(),
+        'ip_address': x[IP_ADDRESS_KEY]
+    } for x in response[INSTANCE_INFORMATION_LIST_KEY]]
+
 
 class AwsService(object):
     """
@@ -56,24 +71,16 @@ class AwsService(object):
         """
         This function will assume that it's running on an EC2 instance with the correct IAM role.
         See https://boto3.amazonaws.com/v1/documentation/api/latest/guide/configuration.html#iam-role for details.
-        :return:
+
+        :return: All visible instances from this instance
         """
-        # local_ssm_client = boto3.client("ssm", region_name=AwsService.region)
         current_instance = AwsInstance()
         local_ssm_client = boto3.client("ssm", region_name=current_instance.get_region())
         try:
             response = local_ssm_client.describe_instance_information()
 
-            return \
-                [
-                    {
-                        'instance_id': x['InstanceId'],
-                        'name': x['ComputerName'],
-                        'os': x['PlatformType'].lower(),
-                        'ip_address': x['IPAddress']
-                    }
-                    for x in response['InstanceInformationList']
-                ]
+            filtered_instances_data = filter_instance_data_from_aws_response(response)
+            return filtered_instances_data
         except botocore.exceptions.ClientError as e:
             print e.response + " " + e.message + " ... " + e.operation_name
             raise e
diff --git a/monkey/common/cloud/test_filter_instance_data_from_aws_response.py b/monkey/common/cloud/test_filter_instance_data_from_aws_response.py
new file mode 100644
index 000000000..04e731f69
--- /dev/null
+++ b/monkey/common/cloud/test_filter_instance_data_from_aws_response.py
@@ -0,0 +1,56 @@
+from unittest import TestCase
+from aws_service import filter_instance_data_from_aws_response
+
+import json
+
+
+class TestFilter_instance_data_from_aws_response(TestCase):
+    def test_filter_instance_data_from_aws_response(self):
+        json_response_full = """
+        {
+            "InstanceInformationList": [ 
+                { 
+                     "ActivationId": "string",
+                     "AgentVersion": "string",
+                     "AssociationOverview": { 
+                        "DetailedStatus": "string",
+                        "InstanceAssociationStatusAggregatedCount": { 
+                           "string" : 6 
+                        }
+                     },
+                     "AssociationStatus": "string",
+                     "ComputerName": "string",
+                     "IamRole": "string",
+                     "InstanceId": "string",
+                     "IPAddress": "string",
+                     "IsLatestVersion": "True",
+                     "LastAssociationExecutionDate": 6,
+                     "LastPingDateTime": 6,
+                     "LastSuccessfulAssociationExecutionDate": 6,
+                     "Name": "string",
+                     "PingStatus": "string",
+                     "PlatformName": "string",
+                     "PlatformType": "string",
+                     "PlatformVersion": "string",
+                     "RegistrationDate": 6,
+                     "ResourceType": "string"
+                }
+            ],
+           "NextToken": "string"
+        }
+        """
+
+        json_response_empty = """
+            {
+                "InstanceInformationList": [],
+                "NextToken": "string"
+            }
+            """
+
+        self.assertEqual(filter_instance_data_from_aws_response(json.loads(json_response_empty)), [])
+        self.assertEqual(
+            filter_instance_data_from_aws_response(json.loads(json_response_full)),
+            [{'instance_id': u'string',
+              'ip_address': u'string',
+              'name': u'string',
+              'os': u'string'}])