diff --git a/libcloudforensics/aws.py b/libcloudforensics/aws.py index 126e73a4..95b5e48b 100644 --- a/libcloudforensics/aws.py +++ b/libcloudforensics/aws.py @@ -22,6 +22,7 @@ import datetime import json import logging +import os import re import boto3 @@ -32,7 +33,10 @@ EC2_SERVICE = 'ec2' ACCOUNT_SERVICE = 'sts' KMS_SERVICE = 'kms' +# Default Amazon Machine Image to use for bootstrapping instances +UBUNTU_1804_AMI = 'ami-0013b3aa57f8a4331' REGEX_TAG_VALUE = re.compile('^.{1,255}$') +STARTUP_SCRIPT = 'scripts/startup.sh' class AWSAccount: @@ -454,6 +458,84 @@ def CreateVolumeFromSnapshot(self, encrypted, name=volume_name) + def GetOrCreateAnalysisVm(self, + vm_name, + boot_volume_size, + ami, + cpu_cores, + packages=None): + """Get or create a new virtual machine for analysis purposes. + + Args: + vm_name (str): The instance name tag of the virtual machine. + boot_volume_size (int): The size of the analysis VM boot volume (in GB). + ami (str): The Amazon Machine Image ID to use to create the VM. + cpu_cores (int): Number of CPU cores for the analysis VM. + packages (list(str)): Optional. List of packages to install in the VM. + + Returns: + tuple(AWSInstance, bool): A tuple with an AWSInstance object and a + boolean indicating if the virtual machine was created (True) or + reused (False). + + Raises: + RuntimeError: If the virtual machine cannot be found or created. + """ + + # Re-use instance if it already exists, or create a new one. + try: + instances = self.GetInstancesByName(vm_name) + if instances: + created = False + return instances[0], created + except RuntimeError: + pass + + instance_type = self._GetInstanceTypeByCPU(cpu_cores) + startup_script = self._ReadStartupScript() + if packages: + startup_script = startup_script.replace('${packages[@]}', ' '.join( + packages)) + + # Install ec2-instance-connect to allow SSH connections from the browser. + startup_script = startup_script.replace( + '(exit ${exit_code})', + 'apt -y install ec2-instance-connect && (exit ${exit_code})') + + client = self.ClientApi(EC2_SERVICE) + # Create the instance in AWS + try: + instance = client.run_instances( + BlockDeviceMappings=[self._GetBootVolumeConfigByAmi( + ami, boot_volume_size)], + ImageId=ami, + MinCount=1, + MaxCount=1, + InstanceType=instance_type, + TagSpecifications=[GetTagForResourceType('instance', vm_name)], + UserData=startup_script, + Placement={'AvailabilityZone': self.default_availability_zone}) + + # If the call to run_instances was successful, then the API response + # contains the instance ID for the new instance. + instance_id = instance['Instances'][0]['InstanceId'] + + # Wait for the instance to be running + client.get_waiter('instance_running').wait(InstanceIds=[instance_id]) + # Wait for the status checks to pass + client.get_waiter('instance_status_ok').wait(InstanceIds=[instance_id]) + + instance = AWSInstance(self, + instance_id, + self.default_region, + self.default_availability_zone, + name=vm_name) + created = True + return instance, created + except client.exceptions.ClientError as exception: + raise RuntimeError('Could not create instance {0:s}: {1:s}'.format( + vm_name, str(exception))) + def GetAccountInformation(self, info): """Get information about the AWS account in use. @@ -591,6 +673,101 @@ def _GenerateVolumeName(self, snapshot, volume_name_prefix=None): return volume_name + def _GetBootVolumeConfigByAmi(self, ami, boot_volume_size): + """Return a boot volume configuration for a given AMI and boot volume size. + + Args: + ami (str): The Amazon Machine Image ID. + boot_volume_size (int): Size of the boot volume, in GB. + + Returns: + dict: A BlockDeviceMappings configuration for the specified AMI. + + Raises: + RuntimeError: If AMI details cannot be found. + """ + + client = self.ClientApi(EC2_SERVICE) + try: + image = client.describe_images(ImageIds=[ami]) + except client.exceptions.ClientError as exception: + raise RuntimeError( + 'Could not find image information for AMI {0:s}: {1:s}'.format( + ami, str(exception))) + + # If the call to describe_images was successful, then the API's response + # is expected to contain at least one image and its corresponding block + # device mappings information. + block_device_mapping = image['Images'][0]['BlockDeviceMappings'][0] + block_device_mapping['Ebs']['VolumeSize'] = boot_volume_size + return block_device_mapping + + @staticmethod + def _GetInstanceTypeByCPU(cpu_cores): + """Return the instance type for the requested number of CPU cores. + + Args: + cpu_cores (int): The number of requested cores. + + Returns: + str: The type of instance that matches the number of cores. + + Raises: + ValueError: If the requested amount of cores is unavailable. + """ + + cpu_cores_to_instance_type = { + 1: 't2.small', + 2: 'm4.large', + 4: 'm4.xlarge', + 8: 'm4.2xlarge', + 16: 'm4.4xlarge', + 32: 'm5.8xlarge', + 40: 'm4.10xlarge', + 48: 'm5.12xlarge', + 64: 'm4.16xlarge', + 96: 'm5.24xlarge', + 128: 'x1.32xlarge' + } + if cpu_cores not in cpu_cores_to_instance_type: + raise ValueError( + 'Cannot start a machine with {0:d} CPU cores. CPU cores should be one' + ' of: {1:s}'.format( + cpu_cores, ', '.join(map(str, cpu_cores_to_instance_type.keys())) + )) + return cpu_cores_to_instance_type[cpu_cores] + + @staticmethod + def _ReadStartupScript(): + """Read and return the startup script that is to be run on the forensics VM. + + Users can either write their own script to install custom packages, + or use the provided one. To use your own script, export a STARTUP_SCRIPT + environment variable with the absolute path to it: + "user@terminal:~$ export STARTUP_SCRIPT='absolute/path/script.sh'" + + Returns: + str: The script to run. + + Raises: + OSError: If the script cannot be opened, read or closed. + """ + + try: + startup_script = os.environ.get('STARTUP_SCRIPT') + if not startup_script: + # Use the provided script + startup_script = os.path.join( + os.path.dirname(os.path.realpath(__file__)), STARTUP_SCRIPT) + startup_script = open(startup_script) + script = startup_script.read() + startup_script.close() + return script + except OSError as exception: + raise OSError( + 'Could not open/read/close the startup script {0:s}: {1:s}'.format( + startup_script, str(exception))) + class AWSInstance: """Class representing an AWS EC2 instance. @@ -649,6 +826,27 @@ def GetBootVolume(self): self.instance_id) raise RuntimeError(error_msg) + def GetVolume(self, volume_id): + """Get a volume attached to the instance by ID. + + Args: + volume_id (str): The ID of the volume to get. + + Returns: + AWSVolume: The AWSVolume object. + + Raises: + RuntimeError: If volume_id is not found amongst the volumes attached + to the instance. + """ + + volume = self.ListVolumes().get(volume_id) + if not volume: + raise RuntimeError( + 'Volume {0:s} is not attached to instance {1:s}'.format( + volume_id, self.instance_id)) + return volume + def ListVolumes(self): """List all volumes for the instance. @@ -661,6 +859,28 @@ def ListVolumes(self): 'Name': 'attachment.instance-id', 'Values': [self.instance_id]}]) + def AttachVolume(self, volume, device_name): + """Attach a volume to the AWS instance. + + Args: + volume (AWSVolume): The AWSVolume object to attach to the instance. + device_name (str): The device name for the volume (e.g. /dev/sdf). + + Raises: + RuntimeError: If the volume could not be attached. + """ + + client = self.aws_account.ClientApi(EC2_SERVICE) + try: + client.attach_volume( + Device=device_name, + InstanceId=self.instance_id, + VolumeId=volume.volume_id) + volume.device_name = device_name + except client.exceptions.ClientError as exception: + raise RuntimeError('Could not attach volume {0:s}: {1:s}'.format( + volume.volume_id, str(exception))) + class AWSElasticBlockStore: """Class representing an AWS EBS resource. @@ -962,12 +1182,60 @@ def CreateVolumeCopy(zone, except RuntimeError as exception: error_msg = 'Copying volume {0:s}: {1!s}'.format( - volume_id, exception) + (volume_id or instance_id), exception) raise RuntimeError(error_msg) return new_volume +def StartAnalysisVm(vm_name, + default_availability_zone, + boot_volume_size, + cpu_cores=4, + ami=UBUNTU_1804_AMI, + attach_volume=None, + device_name=None, + dst_account=None): + """Start a virtual machine for analysis purposes. + + Look for an existing AWS instance with tag name vm_name. If found, + this instance will be started and used as analysis VM. If not found, then a + new vm with that name will be created, started and returned. + + Args: + vm_name (str): The name for the virtual machine. + default_availability_zone (str): Default zone within the region to create + new resources in. + boot_volume_size (int): The size of the analysis VM boot volume (in GB). + cpu_cores (int): Optional. The number of CPU cores to create the machine + with. Default is 4. + ami (str): Optional. The Amazon Machine Image ID to use to create the VM. + Default is a version of Ubuntu 18.04. + attach_volume (AWSVolume): Optional. The volume to attach. + device_name (str): Optional. The name of the device (e.g. /dev/sdf) for the + volume to be attached. Mandatory if attach_volume is provided. + dst_account (str): Optional. The AWS account in which to create the + analysis VM. This is the profile name that is defined in your AWS + credentials file. + + Returns: + tuple(AWSInstance, bool): a tuple with a virtual machine object + and a boolean indicating if the virtual machine was created or not. + + Raises: + RuntimeError: If device_name is missing when attach_volume is provided. + """ + aws_account = AWSAccount(default_availability_zone, aws_profile=dst_account) + analysis_vm, created = aws_account.GetOrCreateAnalysisVm( + vm_name, boot_volume_size, cpu_cores=cpu_cores, ami=ami) + if attach_volume: + if not device_name: + raise RuntimeError('If you want to attach a volume, you must also ' + 'specify a device name for that volume.') + analysis_vm.AttachVolume(attach_volume, device_name) + return analysis_vm, created + + def GetTagForResourceType(resource, name): """Create a dictionary for AWS Tag Specifications. diff --git a/tests/e2e/aws_e2e.py b/tests/e2e/aws_e2e.py index 1a836b93..0b357d68 100644 --- a/tests/e2e/aws_e2e.py +++ b/tests/e2e/aws_e2e.py @@ -59,6 +59,9 @@ def setUpClass(cls): cls.zone = project_info['zone'] cls.volume_to_forensic = project_info.get('volume_id', None) cls.aws = aws.AWSAccount(cls.zone) + cls.analysis_vm_name = 'new-vm-for-analysis' + cls.analysis_vm, _ = aws.StartAnalysisVm( + cls.analysis_vm_name, cls.zone, 10, 4) cls.volumes = [] def test_end_to_end_boot_volume(self): @@ -106,9 +109,45 @@ def test_end_to_end_other_volume(self): self.aws.ResourceApi(EC2_SERVICE).Volume(other_volume_copy.volume_id)) self.assertEqual(self.volumes[-1].volume_id, other_volume_copy.volume_id) + def test_end_to_end_vm(self): + """End to end test on AWS. + + This tests that an analysis VM is correctly created and that a volume + passed to the attach_volume parameter is correctly attached. + """ + + volume_to_attach = aws.CreateVolumeCopy( + self.zone, + volume_id=self.volume_to_forensic) + self.volumes.append(volume_to_attach) + # Create and start the analysis VM and attach the boot volume + self.analysis_vm, _ = aws.StartAnalysisVm( + self.analysis_vm_name, + self.zone, + 10, + 4, + attach_volume=volume_to_attach, + device_name='/dev/sdp' + ) + + # The forensic instance should be live in the analysis AWS account and + # the volume should be attached + instance = self.aws.ResourceApi(EC2_SERVICE).Instance( + self.analysis_vm.instance_id) + self.assertEqual(instance.instance_id, self.analysis_vm.instance_id) + self.assertIn(volume_to_attach.volume_id, + [vol.volume_id for vol in instance.volumes.all()]) + @classmethod def tearDownClass(cls): client = cls.aws.ClientApi(EC2_SERVICE) + # Delete the instance + instance = cls.aws.ResourceApi(EC2_SERVICE).Instance( + cls.analysis_vm.instance_id) + instance.terminate() + client.get_waiter('instance_terminated').wait(InstanceIds=[ + instance.instance_id]) + # Delete the volumes for volume in cls.volumes: log.info('Deleting volume: {0:s}.'.format(volume.volume_id)) diff --git a/tests/unittest/aws_test.py b/tests/unittest/aws_test.py index d7dc7489..e726b3a4 100644 --- a/tests/unittest/aws_test.py +++ b/tests/unittest/aws_test.py @@ -137,6 +137,22 @@ 'UserId': 'fake-user-id' } +MOCK_DESCRIBE_AMI = { + 'Images': [{ + 'BlockDeviceMappings': [{ + 'Ebs': { + 'VolumeSize': None + } + }] + }] +} + +MOCK_RUN_INSTANCES = { + 'Instances': [{ + 'InstanceId': 'new-instance-id' + }] +} + class AWSAccountTest(unittest.TestCase): """Test AWSAccount class.""" @@ -317,6 +333,37 @@ def testCreateVolumeFromSnapshot(self, mock_ec2_api): self.assertEqual( 'prefix-fake-snapshot-d69d57c3-copy', volume_from_snapshot.name) + @mock.patch('libcloudforensics.aws.AWSAccount._ReadStartupScript') + @mock.patch('libcloudforensics.aws.AWSAccount.GetInstancesByName') + @mock.patch('libcloudforensics.aws.AWSAccount.ClientApi') + def testGetOrCreateAnalysisVm(self, + mock_ec2_api, + mock_get_instance, + mock_script): + """Test that a VM is created or retrieved if it already exists.""" + mock_get_instance.return_value = [FAKE_INSTANCE_WITH_NAME] + mock_script.return_value = '' + # GetOrCreateAnalysisVm(vm_name, boot_volume_size, AMI, cpu_cores) where + # vm_name is the name of an analysis instance that already exists. + vm, created = FAKE_AWS_ACCOUNT.GetOrCreateAnalysisVm( + FAKE_INSTANCE_WITH_NAME.name, 1, 'ami-id', 2) + mock_ec2_api.return_value.run_instances.assert_not_called() + self.assertIsInstance(vm, aws.AWSInstance) + self.assertEqual('fake-instance', vm.name) + self.assertFalse(created) + + # GetOrCreateAnalysisVm(non_existing_vm, boot_volume_size, AMI, cpu_cores). + # We mock the GetInstanceById() call to throw a RuntimeError to mimic + # an instance that wasn't found. This should trigger run_instances to be + # called. + mock_get_instance.side_effect = RuntimeError() + vm, created = FAKE_AWS_ACCOUNT.GetOrCreateAnalysisVm( + 'non-existent-instance-name', 1, 'ami-id', 2) + mock_ec2_api.return_value.run_instances.assert_called() + self.assertIsInstance(vm, aws.AWSInstance) + self.assertEqual('non-existent-instance-name', vm.name) + self.assertTrue(created) + @mock.patch('libcloudforensics.aws.AWSAccount.ClientApi') def testGenerateVolumeName(self, mock_ec2_api): """Test the generation of AWS volume name tag. @@ -332,7 +379,27 @@ def testGenerateVolumeName(self, mock_ec2_api): volume_name = FAKE_AWS_ACCOUNT._GenerateVolumeName( FAKE_SNAPSHOT, volume_name_prefix='prefix') + # pylint: enable=protected-access self.assertEqual('prefix-fake-snapshot-d69d57c3-copy', volume_name) + + @mock.patch('libcloudforensics.aws.AWSAccount.ClientApi') + def testGetBootVolumeConfigByAmi(self, mock_ec2_api): + """Test that the boot volume configuration is correctly created.""" + mock_ec2_api.return_value.describe_images.return_value = MOCK_DESCRIBE_AMI + self.assertIsNone( + MOCK_DESCRIBE_AMI['Images'][0]['BlockDeviceMappings'][0]['Ebs']['VolumeSize']) # pylint: disable=line-too-long + # pylint: disable=protected-access + config = FAKE_AWS_ACCOUNT._GetBootVolumeConfigByAmi('ami-id', 50) + # pylint: enable=protected-access + self.assertEqual(50, config['Ebs']['VolumeSize']) + + def testGetInstanceTypeByCPU(self): + """Test that the instance type matches the requested amount of CPU cores.""" + # pylint: disable=protected-access + self.assertEqual('m4.large', FAKE_AWS_ACCOUNT._GetInstanceTypeByCPU(2)) + self.assertEqual('m4.16xlarge', FAKE_AWS_ACCOUNT._GetInstanceTypeByCPU(64)) + self.assertRaises(ValueError, FAKE_AWS_ACCOUNT._GetInstanceTypeByCPU, 0) + self.assertRaises(ValueError, FAKE_AWS_ACCOUNT._GetInstanceTypeByCPU, 256) # pylint: enable=protected-access @@ -378,5 +445,93 @@ def testSnapshot(self, mock_ec2_api): self.assertTrue(snapshot.name.startswith('my-snapshot')) +class AWSTest(unittest.TestCase): + """Test the aws.py public methods.""" + + @mock.patch('libcloudforensics.aws.AWSVolume.Snapshot') + @mock.patch('libcloudforensics.aws.AWSAccount.GetVolumeById') + @mock.patch('libcloudforensics.aws.AWSAccount.GetAccountInformation') + @mock.patch('libcloudforensics.aws.AWSAccount.ClientApi') + def testCreateVolumeCopy1(self, + mock_ec2_api, + mock_account, + mock_get_volume, + mock_snapshot): + """Test that a volume is correctly cloned.""" + FAKE_SNAPSHOT.name = FAKE_VOLUME.volume_id + mock_ec2_api.return_value.create_volume.return_value = MOCK_CREATE_VOLUME + mock_account.return_value = 'fake-account-id' + mock_get_volume.return_value = FAKE_VOLUME + mock_snapshot.return_value = FAKE_SNAPSHOT + + # CreateVolumeCopy(zone, volume_id='fake-volume-id'). This should grab + # the volume 'fake-volume-id'. + new_volume = aws.CreateVolumeCopy( + FAKE_INSTANCE.availability_zone, volume_id=FAKE_VOLUME.volume_id) + mock_get_volume.assert_called_with('fake-volume-id') + self.assertIsInstance(new_volume, aws.AWSVolume) + self.assertTrue(new_volume.name.startswith('evidence-')) + self.assertIn('fake-volume-id', new_volume.name) + self.assertTrue(new_volume.name.endswith('-copy')) + + @mock.patch('libcloudforensics.aws.AWSVolume.Snapshot') + @mock.patch('libcloudforensics.aws.AWSInstance.GetBootVolume') + @mock.patch('libcloudforensics.aws.AWSAccount.GetInstanceById') + @mock.patch('libcloudforensics.aws.AWSAccount.GetAccountInformation') + @mock.patch('libcloudforensics.aws.AWSAccount.ClientApi') + def testCreateVolumeCopy2(self, + mock_ec2_api, + mock_account, + mock_get_instance, + mock_get_volume, + mock_snapshot): + """Test that a volume is correctly cloned.""" + FAKE_SNAPSHOT.name = FAKE_BOOT_VOLUME.volume_id + mock_ec2_api.return_value.create_volume.return_value = MOCK_CREATE_VOLUME + mock_account.return_value = 'fake-account-id' + mock_get_instance.return_value = FAKE_INSTANCE + mock_get_volume.return_value = FAKE_BOOT_VOLUME + mock_snapshot.return_value = FAKE_SNAPSHOT + + # CreateVolumeCopy(zone, instance='fake-instance-id'). This should grab + # the boot volume of the instance. + new_volume = aws.CreateVolumeCopy( + FAKE_INSTANCE.availability_zone, instance_id=FAKE_INSTANCE.instance_id) + mock_get_instance.assert_called_with('fake-instance-id') + self.assertIsInstance(new_volume, aws.AWSVolume) + self.assertTrue(new_volume.name.startswith('evidence-')) + self.assertIn('fake-boot-volume-id', new_volume.name) + self.assertTrue(new_volume.name.endswith('-copy')) + + @mock.patch('libcloudforensics.aws.AWSAccount.ListVolumes') + @mock.patch('libcloudforensics.aws.AWSAccount.ListInstances') + def testCreateVolumeCopy3(self, mock_list_instances, mock_list_volumes): + """Test that a volume is correctly cloned.""" + # Should raise a ValueError exception as no volume_id or instance_id is + # specified. + self.assertRaises( + ValueError, + aws.CreateVolumeCopy, + FAKE_INSTANCE.availability_zone) + + # Should raise a RuntimeError in GetInstanceById as we are querying a + # non-existent instance. + mock_list_instances.return_value = {} + self.assertRaises( + RuntimeError, + aws.CreateVolumeCopy, + FAKE_INSTANCE.availability_zone, + instance_id='non-existent-instance-id') + + # Should raise a RuntimeError in GetVolumeById as we are querying a + # non-existent volume. + mock_list_volumes.return_value = {} + self.assertRaises( + RuntimeError, + aws.CreateVolumeCopy, + FAKE_INSTANCE.availability_zone, + volume_id='non-existent-volume-id') + + if __name__ == '__main__': unittest.main()