Add machine_type param to StartAnalysisVm and pass it through to GetOrCreateAnalysisVm

Author: Fryyyyy

libcloudforensics/providers/gcp/forensics.py

@@ -137,10 +137,11 @@ def StartAnalysisVm(
     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
 

libcloudforensics/providers/gcp/internal/compute.py

@@ -961,6 +961,7 @@ def GetOrCreateAnalysisVm(
       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:

tests/providers/gcp/internal/test_compute.py

@@ -429,6 +429,22 @@ 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')
+
+    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')