def GetInstanceById(self, instance_id: str, region: Optional[str] = None) -> AWSInstance: """Get an instance from an AWS account by its ID. Args: instance_id (str): The instance id. region (str): Optional. The region to look the instance in. If none provided, the default_region associated to the AWSAccount object will be used. Returns: AWSInstance: An Amazon EC2 Instance object. Raises: ResourceNotFoundError: If instance does not exist. """ instances = self.ListInstances(region=region) instance = instances.get(instance_id) if not instance: raise errors.ResourceNotFoundError( 'Instance {0:s} was not found in AWS account'.format(instance_id), __name__) return instance
def testGetOrCreateAnalysisVm(self, mock_vm, mock_nic, mock_instance_type, mock_get_instance, mock_script, mock_ssh_parse): """Test that a VM is created or retrieved if it already exists.""" mock_instance_type.return_value = 'fake-instance-type' mock_nic.return_value = 'fake-network-interface-id' mock_get_instance.return_value = azure_mocks.FAKE_INSTANCE mock_script.return_value = '' mock_ssh_parse.return_value = None vm, created = azure_mocks.FAKE_ACCOUNT.compute.GetOrCreateAnalysisVm( azure_mocks.FAKE_INSTANCE.name, 1, 4, 8192, '') mock_get_instance.assert_called_with(azure_mocks.FAKE_INSTANCE.name) mock_vm.assert_not_called() self.assertIsInstance(vm, compute.AZComputeVirtualMachine) self.assertEqual('fake-vm-name', vm.name) self.assertFalse(created) # We mock the GetInstance() call to throw a ResourceNotFoundError to mimic # an instance that wasn't found. This should trigger a vm to be # created. mock_get_instance.side_effect = errors.ResourceNotFoundError('', __name__) mock_vm.return_value.result.return_value = azure_mocks.MOCK_ANALYSIS_INSTANCE vm, created = azure_mocks.FAKE_ACCOUNT.compute.GetOrCreateAnalysisVm( 'fake-analysis-vm-name', 1, 4, 8192, '') mock_get_instance.assert_called_with('fake-analysis-vm-name') mock_vm.assert_called() self.assertIsInstance(vm, compute.AZComputeVirtualMachine) self.assertEqual('fake-analysis-vm-name', vm.name) self.assertTrue(created)
def GetDisk( self, disk_name: str, resource_group_name: Optional[str] = None) -> 'AZComputeDisk': """Get disk from AZ subscription / resource group. Args: disk_name (str): The disk name. resource_group_name (str): Optional. The resource group name to look the disk in. If none specified, then the disk will be fetched from the AZ subscription. Returns: AZComputeDisk: An Azure Compute Disk object. Raises: ResourceNotFoundError: If the disk was not found in the subscription/ resource group. """ disks = self.ListDisks(resource_group_name=resource_group_name) if disk_name not in disks: raise errors.ResourceNotFoundError( 'Disk {0:s} was not found in subscription {1:s}'.format( disk_name, self.az_account.subscription_id), __name__) return disks[disk_name]
def GetInstance(self, instance_name: str, resource_group_name: Optional[str] = None ) -> 'AZComputeVirtualMachine': """Get instance from AZ subscription / resource group. Args: instance_name (str): The instance name. resource_group_name (str): Optional. The resource group name to look the instance in. If none specified, then the instance will be fetched from the AZ subscription. Returns: AZComputeVirtualMachine: An Azure virtual machine object. Raises: ResourceNotFoundError: If the instance was not found in the subscription/ resource group. """ instances = self.ListInstances(resource_group_name=resource_group_name) if instance_name not in instances: raise errors.ResourceNotFoundError( 'Instance {0:s} was not found in subscription {1:s}'.format( instance_name, self.az_account.subscription_id), __name__) return instances[instance_name]
def GetVolumeById(self, volume_id: str, region: Optional[str] = None) -> AWSVolume: """Get a volume from an AWS account by its ID. Args: volume_id (str): The volume id. region (str): Optional. The region to look the volume in. If none provided, the default_region associated to the AWSAccount object will be used. Returns: AWSVolume: An Amazon EC2 Volume object. Raises: ResourceNotFoundError: If the volume does not exist. """ volumes = self.ListVolumes(region=region) volume = volumes.get(volume_id) if not volume: raise errors.ResourceNotFoundError( 'Volume {0:s} was not found in AWS account'.format(volume_id), __name__) return volume
def Put(self, s3_path: str, filepath: str) -> None: """Upload a local file to an S3 bucket. Keeps the local filename intact. Args: s3_path (str): Path to the target S3 bucket. Ex: s3://test/bucket filepath (str): Path to the file to be uploaded. Ex: /tmp/myfile Raises: ResourceCreationError: If the object couldn't be uploaded. """ client = self.aws_account.ClientApi(common.S3_SERVICE) if not s3_path.startswith('s3://'): s3_path = 's3://' + s3_path if not s3_path.endswith('/'): s3_path = s3_path + '/' try: (bucket, path) = gcp_storage.SplitStoragePath(s3_path) client.upload_file( filepath, bucket, '{0:s}{1:s}'.format(path, os.path.basename(filepath))) except FileNotFoundError as exception: raise errors.ResourceNotFoundError( 'Could not upload file {0:s}: {1:s}'.format( filepath, str(exception)), __name__) from exception except client.exceptions.ClientError as exception: raise errors.ResourceCreationError( 'Could not upload file {0:s}: {1:s}'.format( filepath, str(exception)), __name__) from exception
def testGetOrCreateAnalysisVm(self, mock_gce_api, mock_block_operation, mock_get_instance): """Test that a new virtual machine is created if it doesn't exist, or that the existing one is returned. """ images = mock_gce_api.return_value.images images.return_value.getFromFamily.return_value.execute.return_value = { 'selfLink': 'value' } instances = mock_gce_api.return_value.instances instances.return_value.insert.return_value.execute.return_value = None mock_block_operation.return_value = None mock_get_instance.return_value = gcp_mocks.FAKE_ANALYSIS_VM # GetOrCreateAnalysisVm(existing_vm, boot_disk_size) vm, created = gcp_mocks.FAKE_ANALYSIS_PROJECT.compute.GetOrCreateAnalysisVm( gcp_mocks.FAKE_ANALYSIS_VM.name, boot_disk_size=1) mock_get_instance.assert_called_with(gcp_mocks.FAKE_ANALYSIS_VM.name) instances.return_value.insert.assert_not_called() self.assertIsInstance(vm, compute.GoogleComputeInstance) self.assertEqual('fake-analysis-vm', vm.name) self.assertFalse(created) # GetOrCreateAnalysisVm(non_existing_vm, boot_disk_size) mocking the # GetInstance() call to throw a ResourceNotFoundError error to mimic an # instance that wasn't found mock_get_instance.side_effect = errors.ResourceNotFoundError( '', __name__) vm, created = gcp_mocks.FAKE_ANALYSIS_PROJECT.compute.GetOrCreateAnalysisVm( 'non-existent-analysis-vm', boot_disk_size=1) mock_get_instance.assert_called_with('non-existent-analysis-vm') instances.return_value.insert.assert_called() self.assertIsInstance(vm, compute.GoogleComputeInstance) self.assertEqual('non-existent-analysis-vm', vm.name) self.assertTrue(created)
def GetBootDisk(self) -> 'GoogleComputeDisk': """Get the virtual machine boot disk. Returns: GoogleComputeDisk: Disk object. Raises: ResourceNotFoundError: If no boot disk could be found. """ for disk in self.GetValue('disks'): if disk['boot']: disk_name = disk['source'].split('/')[-1] return GoogleCloudCompute(self.project_id).GetDisk(disk_name=disk_name) raise errors.ResourceNotFoundError( 'Boot disk not found for instance {0:s}'.format(self.name), __name__)
def GetDisk(self, disk_name: str) -> 'AZComputeDisk': """Get a disk attached to the instance by ID. Args: disk_name (str): The ID of the disk to get. Returns: AZComputeDisk: The disk object. Raises: ResourceNotFoundError: If disk_name is not found amongst the disks attached to the instance. """ disks = self.ListDisks() if disk_name not in disks: raise errors.ResourceNotFoundError( 'Disk {0:s} was not found in instance {1:s}'.format( disk_name, self.resource_id), __name__) return disks[disk_name]
def _GetBootVolumeConfigByAmi(self, ami: str, boot_volume_size: int, boot_volume_type: str) -> Dict[str, Any]: """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. boot_volume_type (str): Type of the boot volume. Returns: Dict[str, str|Dict]]: A BlockDeviceMappings configuration for the specified AMI. Raises: ResourceNotFoundError: If AMI details cannot be found. """ client = self.aws_account.ClientApi(common.EC2_SERVICE) try: image = client.describe_images(ImageIds=[ami]) except client.exceptions.ClientError as exception: raise errors.ResourceNotFoundError( 'Could not find image information for AMI {0:s}: {1!s}'.format( ami, exception), __name__) # 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. # pylint: disable=line-too-long block_device_mapping = image['Images'][0]['BlockDeviceMappings'][0] # type: Dict[str, Any] # pylint: enable=line-too-long block_device_mapping['Ebs']['VolumeSize'] = boot_volume_size block_device_mapping['Ebs']['VolumeType'] = boot_volume_type if boot_volume_type == 'io1': # If using the io1 volume type, we must specify Iops, see # https://boto3.amazonaws.com/v1/documentation/api/latest/reference/ # services/ec2.html#EC2.Client.create_volume. io1 volumes allow for a # ratio of 50 IOPS per 1 GiB. block_device_mapping['Ebs']['Iops'] = boot_volume_size * 50 return block_device_mapping
def GetDisk(self, disk_name: str) -> 'GoogleComputeDisk': """Gets a disk attached to this virtual machine disk by name. Args: disk_name (str): The name of the disk to get. Returns: GoogleComputeDisk: Disk object. Raises: ResourceNotFoundError: If disk name is not found among those attached to the instance. """ for disk in self.GetValue('disks'): if disk['source'].split('/')[-1] == disk_name: return GoogleCloudCompute(self.project_id).GetDisk(disk_name=disk_name) raise errors.ResourceNotFoundError( 'Disk {0:s} was not found in instance {1:s}'.format( disk_name, self.name), __name__)
def GetDisk(self, disk_name: str) -> 'GoogleComputeDisk': """Get a GCP disk object. Args: disk_name (str): Name of the disk. Returns: GoogleComputeDisk: Disk object. Raises: ResourceNotFoundError: When the specified disk cannot be found in project. """ disks = self.Disks() disk = disks.get(disk_name) if not disk: raise errors.ResourceNotFoundError( 'Disk {0:s} was not found in project {1:s}'.format( disk_name, self.project_id), __name__) return disk
def GetInstance(self, instance_name: str) -> 'GoogleComputeInstance': """Get instance from project. Args: instance_name (str): The instance name. Returns: GoogleComputeInstance: A Google Compute Instance object. Raises: ResourceNotFoundError: If instance does not exist. """ instances = self.Instances() instance = instances.get(instance_name) if not instance: raise errors.ResourceNotFoundError( 'Instance {0:s} was not found in project {1:s}'.format( instance_name, self.project_id), __name__) return instance
def GetBootDisk(self) -> 'AZComputeDisk': """Get the instance's boot disk. Returns: AZComputeDisk: Disk object if the disk is found. Raises: ResourceNotFoundError: If no boot disk could be found. """ # pylint: disable=line-too-long disks = self.az_account.compute.ListDisks( resource_group_name=self.resource_group_name) # type: Dict[str, AZComputeDisk] # pylint: enable=line-too-long boot_disk_name = self.compute_client.virtual_machines.get( self.resource_group_name, self.name).storage_profile.os_disk.name if boot_disk_name not in disks: raise errors.ResourceNotFoundError( 'Boot disk not found for instance {0:s}'.format(self.resource_id), __name__) return disks[boot_disk_name]
def GetVolume(self, volume_id: str) -> 'ebs.AWSVolume': """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: ResourceNotFoundError: If volume_id is not found amongst the volumes attached to the instance. """ volume = self.ListVolumes().get(volume_id) if not volume: raise errors.ResourceNotFoundError( 'Volume {0:s} is not attached to instance {1:s}'.format( volume_id, self.instance_id), __name__) return volume
def GetBootVolume(self) -> 'ebs.AWSVolume': """Get the instance's boot volume. Returns: AWSVolume: Volume object if the volume is found. Raises: ResourceNotFoundError: If no boot volume could be found. """ boot_device = self.aws_account.ResourceApi( common.EC2_SERVICE).Instance(self.instance_id).root_device_name volumes = self.ListVolumes() for volume_id in volumes: if volumes[volume_id].device_name == boot_device: return volumes[volume_id] raise errors.ResourceNotFoundError( 'Boot volume not found for instance: {0:s}'.format(self.instance_id), __name__)
def CreateDiskCopy( src_proj: str, dst_proj: str, zone: str, instance_name: Optional[str] = None, disk_name: Optional[str] = None, disk_type: Optional[str] = None) -> 'compute.GoogleComputeDisk': """Creates a copy of a Google Compute Disk. Args: src_proj (str): Name of project that holds the disk to be copied. dst_proj (str): Name of project to put the copied disk in. zone (str): Zone where the new disk is to be created. instance_name (str): Optional. Instance using the disk to be copied. disk_name (str): Optional. Name of the disk to copy. If None, instance_name must be specified and the boot disk will be copied. disk_type (str): Optional. URL of the disk type resource describing which disk type to use to create the disk. The default behavior is to use the same disk type as the source disk. Returns: GoogleComputeDisk: A Google Compute Disk object. Raises: ResourceNotFoundError: If the GCP resource is not found. CredentialsConfigurationError: If the library could not authenticate to GCP. RuntimeError: If an unknown HttpError is thrown. ResourceCreationError: If there are errors copying the disk. ValueError: If both instance_name and disk_name are missing. """ if not instance_name and not disk_name: raise ValueError( 'You must specify at least one of [instance_name, disk_name].') src_project = gcp_project.GoogleCloudProject(src_proj) dst_project = gcp_project.GoogleCloudProject(dst_proj, default_zone=zone) try: if disk_name: disk_to_copy = src_project.compute.GetDisk(disk_name) elif instance_name: instance = src_project.compute.GetInstance(instance_name) disk_to_copy = instance.GetBootDisk() if not disk_type: disk_type = disk_to_copy.GetDiskType() logger.info('Disk copy of {0:s} started...'.format( disk_to_copy.name)) snapshot = disk_to_copy.Snapshot() logger.debug('Snapshot created: {0:s}'.format(snapshot.name)) new_disk = dst_project.compute.CreateDiskFromSnapshot( snapshot, disk_name_prefix='evidence', disk_type=disk_type) logger.info( 'Disk {0:s} successfully copied to {1:s}'.format( disk_to_copy.name, new_disk.name)) snapshot.Delete() logger.debug('Snapshot {0:s} deleted.'.format(snapshot.name)) except (RefreshError, DefaultCredentialsError) as exception: raise errors.CredentialsConfigurationError( 'Something is wrong with your Application Default Credentials. Try ' 'running: $ gcloud auth application-default login: {0!s}'.format( exception), __name__) from exception except HttpError as exception: if exception.resp.status == 403: raise errors.CredentialsConfigurationError( 'Make sure you have the appropriate permissions on the project', __name__) from exception if exception.resp.status == 404: raise errors.ResourceNotFoundError( 'GCP resource not found. Maybe a typo in the project / instance / ' 'disk name?', __name__) from exception raise RuntimeError(exception) from exception except RuntimeError as exception: raise errors.ResourceCreationError( 'Cannot copy disk "{0:s}": {1!s}'.format(disk_name, exception), __name__) from exception return new_disk