diff --git a/libcloudforensics/providers/gcp/forensics.py b/libcloudforensics/providers/gcp/forensics.py index 848dcd41..b2892647 100644 --- a/libcloudforensics/providers/gcp/forensics.py +++ b/libcloudforensics/providers/gcp/forensics.py @@ -130,17 +130,18 @@ def CreateDiskCopy( return new_disk -def StartAnalysisVm( +def StartAnalysisVm( # pylint: disable=too-many-arguments project: str, vm_name: str, zone: str, boot_disk_size: int = 10, boot_disk_type: str = 'pd-standard', cpu_cores: int = 4, + machine_type: Optional[str] = None, attach_disks: Optional[List[str]] = None, image_project: str = 'ubuntu-os-cloud', image_family: str = 'ubuntu-2204-lts', - packages: Optional[List[str]] = None + packages: Optional[List[str]] = None, ) -> Tuple['compute.GoogleComputeInstance', bool]: """Start a virtual machine for analysis purposes. @@ -149,18 +150,18 @@ def StartAnalysisVm( vm_name: The name of the virtual machine. zone: Zone for the virtual machine. boot_disk_size: The size of the analysis VM boot disk (in GB). - boot_disk_type: URL of the disk type resource describing - which disk type to use to create the disk. Use pd-standard for a - standard disk and pd-ssd for a SSD disk. + boot_disk_type: URL of the disk type resource describing which disk type to + use to create the disk. Use pd-standard for a standard disk and pd-ssd for + a SSD disk. cpu_cores: The number of CPU cores to create the machine with. - attach_disks: List of disk names to attach. Default - behaviour is to search in zonal disks then regional disks, when using - regional disks CreateInstanceFromArguments from GoogleCloudCompute is - recommended to avoid name colisions with zonal disks. - image_project: Name of the project where the analysis VM - image is hosted. - image_family: Name of the image to use to create the - analysis VM. + machine_type: Machine type for the virtual machine. If specified, cpu_cores + will be ignored. + attach_disks: List of disk names to attach. Default behaviour is to search + in zonal disks then regional disks, when using regional disks + CreateInstanceFromArguments from GoogleCloudCompute is recommended to + avoid name colisions with zonal disks. + image_project: Name of the project where the analysis VM image is hosted. + image_family: Name of the image to use to create the analysis VM. packages: List of extra packages to install in the VM. Returns: @@ -178,10 +179,17 @@ def StartAnalysisVm( disk = proj.compute.GetRegionDisk(disk_name) data_disks.append(disk) analysis_vm, created = proj.compute.GetOrCreateAnalysisVm( - vm_name, boot_disk_size, disk_type=boot_disk_type, cpu_cores=cpu_cores, - image_project=image_project, image_family=image_family, + vm_name, + boot_disk_size, + disk_type=boot_disk_type, + cpu_cores=cpu_cores, + machine_type=machine_type, + image_project=image_project, + image_family=image_family, packages=packages, - data_disks=data_disks, zone=zone) + data_disks=data_disks, + zone=zone, + ) logger.info('VM started.') return analysis_vm, created diff --git a/libcloudforensics/providers/gcp/internal/compute.py b/libcloudforensics/providers/gcp/internal/compute.py index e4158160..4b005a06 100644 --- a/libcloudforensics/providers/gcp/internal/compute.py +++ b/libcloudforensics/providers/gcp/internal/compute.py @@ -955,12 +955,13 @@ def CreateInstanceFromArguments( # pylint: disable=too-many-arguments,too-many- pass return self.CreateInstanceFromRequest(request_body, compute_zone) - def GetOrCreateAnalysisVm( + def GetOrCreateAnalysisVm( # pylint: disable=too-many-arguments self, vm_name: str, boot_disk_size: int = 10, disk_type: str = 'pd-standard', cpu_cores: int = 4, + machine_type: Optional[str] = None, image_project: str = 'ubuntu-os-cloud', image_family: str = 'ubuntu-2204-lts', packages: Optional[List[str]] = None, @@ -984,6 +985,8 @@ def GetOrCreateAnalysisVm( which disk type to use to create the disk. Default is pd-standard. Use pd-ssd to have a SSD disk. cpu_cores: Number of CPU cores for the virtual machine. + machine_type: Machine type for the virtual machine. If specified, + cpu_cores will be ignored. image_project: Name of the project where the analysis VM image is hosted. image_family: Name of the image to use to create the @@ -1022,11 +1025,13 @@ def GetOrCreateAnalysisVm( pass compute_zone = zone if zone else self.default_zone - if cpu_cores not in E2_STANDARD_CPU_CORES: - raise ValueError( - 'Number of requested CPU cores ({0:d}) not available for machine type' - ' {1:s}'.format(cpu_cores, DEFAULT_MACHINE_TYPE)) - machine_type = '{0:s}-{1:d}'.format(DEFAULT_MACHINE_TYPE, cpu_cores) + if machine_type is None: + if cpu_cores not in E2_STANDARD_CPU_CORES: + raise ValueError( + 'Number of requested CPU cores ({0:d}) ' + 'not available for machine type' + ' {1:s}'.format(cpu_cores, DEFAULT_MACHINE_TYPE)) + machine_type = '{0:s}-{1:d}'.format(DEFAULT_MACHINE_TYPE, cpu_cores) startup_script = utils.ReadStartupScript(utils.FORENSICS_STARTUP_SCRIPT_GCP) if packages: diff --git a/tests/providers/gcp/internal/test_compute.py b/tests/providers/gcp/internal/test_compute.py index 53c91383..3b2c21d0 100644 --- a/tests/providers/gcp/internal/test_compute.py +++ b/tests/providers/gcp/internal/test_compute.py @@ -429,6 +429,24 @@ def testGetOrCreateAnalysisVmExist(self, mock_create_from_args, mock_get_instanc self.assertFalse(created) mock_create_from_args.assert_not_called() + @typing.no_type_check + @mock.patch('libcloudforensics.providers.gcp.internal.compute.GoogleCloudCompute.GetInstance') + @mock.patch('libcloudforensics.providers.gcp.internal.compute.GoogleCloudCompute.CreateInstanceFromArguments') + def testGetOrCreateAnalysisVmWithMachineType(self, mock_create_from_args, mock_get_instance): + """Test creating analysis VM with specific machine type.""" + mock_get_instance.side_effect = errors.ResourceNotFoundError('', __name__) + mock_create_from_args.return_value = gcp_mocks.FAKE_ANALYSIS_VM + + vm, created = gcp_mocks.FAKE_ANALYSIS_PROJECT.compute.GetOrCreateAnalysisVm( + gcp_mocks.FAKE_ANALYSIS_VM.name, machine_type='custom-machine-type') + + mock_get_instance.assert_called_with(gcp_mocks.FAKE_ANALYSIS_VM.name) + self.assertIsInstance(vm, compute.GoogleComputeInstance) + self.assertTrue(created) + mock_create_from_args.assert_called_once() + args, _ = mock_create_from_args.call_args + self.assertEqual('custom-machine-type', args[1]) + @typing.no_type_check @mock.patch('libcloudforensics.providers.gcp.internal.compute.GoogleCloudCompute.ListInstanceByLabels') @mock.patch('libcloudforensics.providers.gcp.internal.common.GoogleCloudComputeClient.GceApi')