Exemplo n.º 1
0
def AMI_builder(
    AWS_access_key_id,
    AWS_secret_access_key,
    region_name,
    base_image_id,
    os,
    security_group_id,
    AMI_name,
    RPM_package_version,
    APT_OSS_version,
):
    """
    Builds the ODFE AMI
    args: 
        AWS_access_key_id: str, aws key id
        AWS_secret_access_key: str, aws secret access key
        region_name: str, region where the Instance will be created
        base_image_id: str, base os AMI id,  
        os: str, ubuntu or amazonLinux
        security_group_id: str, security group with port 22 open
        AMI_name: str, Name of the AMI that will be created
        RPM_package_version: str, version of ODFE to be installed if RPM is used for installation(used in amazon linux)
        APT_OSS_version: str, version of Elasticsearch OSS to be installed if apt is used for installation(used in ubuntu)
    returns: none
    """
    try:
        instance = Instance(
            AWS_access_key_id=AWS_access_key_id,
            AWS_secret_access_key=AWS_secret_access_key,
            region_name=region_name,
            base_image_id=base_image_id,
            os=os,  # ubuntu, amazonLinux
            security_group_id=security_group_id,
            AMI_name=AMI_name,
            RPM_package_version=RPM_package_version,
            APT_OSS_version=APT_OSS_version,
        )
    except Exception as err:
        logging.error("Could not bring up the instance. " + str(err))
        sys.exit(-1)
    AMI_id = ""
    installation_failed = False
    try:
        instance.wait_until_ready()
    except Exception as err:
        logging.error("Could not bring the instance to ready state. " +
                      str(err))
        installation_failed = True
    else:
        try:
            instance.install_ODFE()
            AMI_id = instance.create_AMI()
        except Exception as err:
            installation_failed = True
            logging.error(
                "AMI creation failed there was an error see the logs. " +
                str(err))
    finally:
        try:
            instance.cleanup_instance()
        except Exception as err:
            logging.error(
                "Could not cleanup the instance. There could be an instance currently running, terminate it. "
                + str(err))
            installation_failed = True
    if installation_failed:
        sys.exit(-1)
    # copy the AMI to the required regions
    ec2_client = boto3.client(
        "ec2",
        aws_access_key_id=AWS_access_key_id,
        aws_secret_access_key=AWS_secret_access_key,
        region_name=region_name,
    )
    AMI_copy_regions = [
        region["RegionName"]
        for region in ec2_client.describe_regions()["Regions"]
    ]
    AMI_copy_regions.remove(region_name)  # since AMI is created here
    copy_AMI_to_regions(
        AWS_access_key_id=AWS_access_key_id,
        AWS_secret_access_key=AWS_secret_access_key,
        AMI_id=AMI_id,
        AMI_name=AMI_name,
        AMI_source_region=region_name,
        AMI_copy_regions=AMI_copy_regions,
    )
class TestInstance(unittest.TestCase):
    @patch("lib.instance.boto3")
    @patch("lib.instance.Instance._create_key_pair")
    def setUp(self, create_key_pair, boto3_mock):
        """
        Initialise mock instance for the test cases
        """
        create_key_pair.return_value = "fakekey.pem"
        self.instance = Instance(
            AWS_access_key_id="1234",
            AWS_secret_access_key="1234",
            region_name="abc_region",
            base_image_id="ami-123",
            os="amazonLinux",
            security_group_id="sg-1234",
            AMI_name="OpenDistroAMI",
            RPM_package_version="1.0.0",
            APT_OSS_version="1.0.0",
        )
        self.boto_instance = boto3_mock.resource().create_instances(
        )[0]  # i could not figure out the any other way to mock boto3 instance
        self.instance.key_pair_name = "ODFEAMIInstanceKey"

    def test_initialization(self):
        """
        Fail with descriptive error if required initialization have failed
        """
        self.assertEqual("amazonLinux", self.instance.os)
        self.assertEqual("OpenDistroAMI", self.instance.AMI_name)
        self.assertEqual("1.0.0", self.instance.RPM_package_version)
        self.assertEqual("1.0.0", self.instance.APT_OSS_version)

    def test_incorrect_region(self):
        """
        Fail with descriptive error if Instance can be created for wrong region
        """
        with self.assertRaises(ValueError):
            Instance(
                region_name="wrong_region",
                AWS_access_key_id="1234",
                AWS_secret_access_key="1234",
                base_image_id="ami-123",
                os="amazonLinux",
                security_group_id="sg-1234",
                AMI_name="OpenDistroAMI",
                RPM_package_version="1.0.0",
                APT_OSS_version="1.0.0",
            )

    def test_incorrect_credentials(self):
        """
        Fail with descriptive error if Instance can be created with wrong credentials
        """
        with self.assertRaises(ClientError):
            Instance(
                AWS_access_key_id="wrong",
                AWS_secret_access_key="wrong",
                region_name="us-east-2",
                base_image_id="ami-123",
                os="amazonLinux",
                security_group_id="sg-1234",
                AMI_name="OpenDistroAMI",
                RPM_package_version="1.0.0",
                APT_OSS_version="1.0.0",
            )

    def test_wait_until_ready(self):
        """
        Exercises the code path for waiting until instance is ready to be used
        """
        self.boto_instance.wait_until_running.side_effect = [
            "",
            WaiterError("blah", "blah", "blah"),
        ]
        self.boto_instance.state = -1
        self.instance.wait_until_ready()
        self.boto_instance.wait_until_running.assert_called()
        self.assertRaises(WaiterError, self.instance.wait_until_ready)

    @patch("lib.instance.Instance.__init__")
    def test_create_key_pair(self, init_mock):
        """
        Exercises the code path for creating key pair
        """
        init_mock.return_value = None
        instance = Instance()
        instance.key_pair_name = self.instance.key_pair_name
        instance.ec2_client = Mock()
        instance.ec2_client.create_key_pair.return_value = {
            "KeyMaterial": "dummy key"
        }
        key_path = instance._create_key_pair()
        self.assertEquals(key_path, f"./{self.instance.key_pair_name}.pem")
        key = open(key_path, "r")
        self.assertEquals(key.read(), "dummy key")
        os.remove(f"./{self.instance.key_pair_name}.pem")

    def test_create_AMI(self):
        """
        Exercises the code path for creating AMI
        """
        self.boto_instance.create_image.return_value = Mock(image_id="123")
        self.instance.create_AMI()
        self.assertEquals(self.instance.snapshot.image_id, "123")
        self.instance.ec2_client.get_waiter().wait.side_effect = WaiterError(
            "bla", "blah", "blah")
        self.assertRaises(WaiterError, self.instance.create_AMI)

    @patch("lib.instance.ODFEInstaller")
    def test_install_ODFE(self, mock_installer):
        """
        Exercises the code path for open distro installation in the instance
        """
        self.instance.install_ODFE()
        mock_installer().install.assert_called()
        mock_installer().install.side_effect = Exception
        self.assertRaises(Exception, self.instance.install_ODFE)

    def test_cleanup_instance(self):
        """
        Exercises the code path for cleaning up the instance after use
        """
        self.boto_instance.wait_until_terminated.side_effect = [
            "",
            WaiterError("blah", "blah", "blah"),
            "",
        ]
        temp = open(f"./{self.instance.key_pair_name}.pem", "w")
        self.instance.key_path = f"./{self.instance.key_pair_name}.pem"
        self.instance.cleanup_instance()
        self.boto_instance.wait_until_terminated.assert_called()
        temp = open(f"./{self.instance.key_pair_name}.pem", "w")
        self.assertRaises(WaiterError, self.instance.cleanup_instance)
        os.remove(f"./{self.instance.key_pair_name}.pem")
        self.assertRaises(FileNotFoundError, self.instance.cleanup_instance)