def testReadStartupScript(self): """Test that the startup script is correctly read.""" # No environment variable set, reading default script # pylint: disable=protected-access script = utils.ReadStartupScript() self.assertTrue(script.startswith('#!/bin/bash')) self.assertTrue(script.endswith('(exit ${exit_code})\n')) # Environment variable set to custom script os.environ['STARTUP_SCRIPT'] = os.path.join( os.path.dirname( os.path.dirname(os.path.dirname(os.path.realpath(__file__)))), STARTUP_SCRIPT) script = utils.ReadStartupScript() self.assertEqual('# THIS IS A CUSTOM BASH SCRIPT', script) # Bogus environment variable, should raise an exception os.environ['STARTUP_SCRIPT'] = '/bogus/path' with self.assertRaises(OSError): utils.ReadStartupScript() os.environ['STARTUP_SCRIPT'] = ''
def GetOrCreateAnalysisVm( self, vm_name: str, boot_volume_size: int, ami: str, cpu_cores: int, packages: Optional[List[str]] = None, ssh_key_name: Optional[str] = None ) -> Tuple[ec2.AWSInstance, bool]: """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. ssh_key_name (str): Optional. A SSH key pair name linked to the AWS account to associate with the VM. If none provided, the VM can only be accessed through in-browser SSH from the AWS management console with the EC2 client connection package (ec2-instance-connect). Note that if this package fails to install on the target VM, then the VM will not be accessible. It is therefore recommended to fill in this parameter. 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 = common.GetInstanceTypeByCPU(cpu_cores) startup_script = utils.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(common.EC2_SERVICE) vm_args = { 'BlockDeviceMappings': [self._GetBootVolumeConfigByAmi(ami, boot_volume_size)], 'ImageId': ami, 'MinCount': 1, 'MaxCount': 1, 'InstanceType': instance_type, 'TagSpecifications': [common.GetTagForResourceType('instance', vm_name)], 'UserData': startup_script, 'Placement': { 'AvailabilityZone': self.default_availability_zone } } if ssh_key_name: vm_args['KeyName'] = ssh_key_name # Create the instance in AWS try: instance = client.run_instances(**vm_args) # 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]) except (client.exceptions.ClientError, botocore.exceptions.WaiterError) as exception: raise RuntimeError('Could not create instance {0:s}: {1:s}'.format( vm_name, str(exception))) instance = ec2.AWSInstance(self, instance_id, self.default_region, self.default_availability_zone, name=vm_name) created = True return instance, created
def GetOrCreateAnalysisVm( self, vm_name: str, boot_disk_size: int, cpu_cores: int, memory_in_mb: int, ssh_public_key: str, region: Optional[str] = None, packages: Optional[List[str]] = None, tags: Optional[Dict[str, str]] = None ) -> Tuple['AZComputeVirtualMachine', bool]: """Get or create a new virtual machine for analysis purposes. Args: vm_name (str): The instance name tag of the virtual machine. boot_disk_size (int): The size of the analysis VM boot volume (in GB). cpu_cores (int): Number of CPU cores for the analysis VM. memory_in_mb (int): The memory size (in MB) for the analysis VM. ssh_public_key (str): A SSH public key data to associate with the VM. This must be provided as otherwise the VM will not be accessible. region (str): Optional. The region in which to create the vm. If not provided, the vm will be created in the default_region associated to the AZAccount object. packages (List[str]): Optional. List of packages to install in the VM. tags (Dict[str, str]): Optional. A dictionary of tags to add to the instance, for example {'TicketID': 'xxx'}. An entry for the instance name is added by default. Returns: Tuple[AZComputeVirtualMachine, bool]: A tuple with an AZComputeVirtualMachine 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: instance = self.GetInstance(vm_name) if instance: created = False return instance, created except RuntimeError: pass # Validate SSH public key format try: sshpubkeys.SSHKey(ssh_public_key, strict=True).parse() except sshpubkeys.InvalidKeyError as exception: raise RuntimeError('The provided public SSH key is invalid: ' '{0:s}'.format(str(exception))) instance_type = self._GetInstanceType(cpu_cores, memory_in_mb) startup_script = utils.ReadStartupScript() if packages: startup_script = startup_script.replace('${packages[@]}', ' '.join(packages)) if not region: region = self.az_account.default_region creation_data = { 'location': region, 'properties': { 'hardwareProfile': { 'vmSize': instance_type }, 'storageProfile': { 'imageReference': { 'sku': common.UBUNTU_1804_SKU, 'publisher': 'Canonical', 'version': 'latest', 'offer': 'UbuntuServer' } }, 'osDisk': { 'caching': 'ReadWrite', 'managedDisk': { 'storageAccountType': 'Standard_LRS' }, 'name': 'os-disk-{0:s}'.format(vm_name), 'diskSizeGb': boot_disk_size, 'createOption': models.DiskCreateOption.from_image }, 'osProfile': { 'adminUsername': '******', 'computerName': vm_name, # Azure requires the startup script to be sent as a b64 string 'customData': base64.b64encode( str.encode(startup_script)).decode('utf-8'), 'linuxConfiguration': { 'ssh': { 'publicKeys': [{ 'path': '/home/AzureUser/.ssh/authorized_keys', 'keyData': ssh_public_key }] } } }, 'networkProfile': { 'networkInterfaces': [ # pylint: disable=line-too-long # This is necessary when creating a VM from the SDK. # See https://docs.microsoft.com/en-us/azure/virtual-machines/windows/python # pylint: enable=line-too-long { 'id': self.az_account.network.CreateNetworkInterface( vm_name, region) } ] } } } # type: Dict[str, Any] if tags: creation_data['tags'] = tags try: request = self.compute_client.virtual_machines.create_or_update( self.az_account.default_resource_group_name, vm_name, creation_data) while not request.done(): sleep(5) # Wait 5 seconds before checking disk status again vm = request.result() except azure_exceptions.CloudError as exception: raise RuntimeError('Could not create instance {0:s}: {1:s}'.format( vm_name, str(exception))) instance = AZComputeVirtualMachine(self.az_account, vm.id, vm.name, vm.location, zones=vm.zones) created = True return instance, created
def GetOrCreateAnalysisVm( self, vm_name: str, boot_disk_size: int, disk_type: str = 'pd-standard', cpu_cores: int = 4, image_project: str = 'ubuntu-os-cloud', image_family: str = 'ubuntu-1804-lts', # pylint: disable=line-too-long packages: Optional[List[str]] = None ) -> Tuple['GoogleComputeInstance', bool]: # pylint: enable=line-too-long """Get or create a new virtual machine for analysis purposes. If none of the optional parameters are specified, then by default the analysis VM that will be created will run Ubuntu 18.04 LTS. A default set of forensic tools is also installed (a custom one may be provided using the 'packages' argument). Args: vm_name (str): Name of the virtual machine. boot_disk_size (int): The size of the analysis VM boot disk (in GB). disk_type (str): Optional. URL of the disk type resource describing which disk type to use to create the disk. Default is pd-standard. Use pd-ssd to have a SSD disk. cpu_cores (int): Optional. Number of CPU cores for the virtual machine. image_project (str): Optional. Name of the project where the analysis VM image is hosted. image_family (str): Optional. Name of the image to use to create the analysis VM. packages (List[str]): Optional. List of packages to install in the VM. Returns: Tuple(GoogleComputeInstance, bool): A tuple with a virtual machine object and a boolean indicating if the virtual machine was created or not. Raises: RuntimeError: If virtual machine cannot be created. """ if not self.default_zone: raise RuntimeError('Cannot create VM, zone information is missing') # Re-use instance if it already exists, or create a new one. try: instance = self.GetInstance(vm_name) created = False return instance, created except RuntimeError: pass machine_type = 'zones/{0}/machineTypes/n1-standard-{1:d}'.format( self.default_zone, cpu_cores) ubuntu_image = self.GceApi().images().getFromFamily( project=image_project, family=image_family).execute() source_disk_image = ubuntu_image['selfLink'] startup_script = utils.ReadStartupScript() if packages: startup_script = startup_script.replace('${packages[@]}', ' '.join(packages)) config = { 'name': vm_name, 'machineType': machine_type, 'disks': [{ 'boot': True, 'autoDelete': True, 'initializeParams': { 'diskType': 'projects/{0:s}/zones/{1:s}/diskTypes/{2:s}'.format( self.project_id, self.default_zone, disk_type), 'sourceImage': source_disk_image, 'diskSizeGb': boot_disk_size, } }], 'networkInterfaces': [{ 'network': 'global/networks/default', 'accessConfigs': [{ 'type': 'ONE_TO_ONE_NAT', 'name': 'External NAT' }] }], 'serviceAccounts': [{ 'email': 'default', 'scopes': [ 'https://www.googleapis.com/auth/devstorage.read_write', 'https://www.googleapis.com/auth/logging.write' ] }], 'metadata': { 'items': [{ 'key': 'startup-script', # Analysis software to install. 'value': startup_script }] } } gce_instance_client = self.GceApi().instances() request = gce_instance_client.insert(project=self.project_id, zone=self.default_zone, body=config) response = request.execute() self.BlockOperation(response, zone=self.default_zone) instance = GoogleComputeInstance(project_id=self.project_id, zone=self.default_zone, name=vm_name) created = True return instance, created