diff --git a/libcloudforensics/providers/gcp/internal/compute.py b/libcloudforensics/providers/gcp/internal/compute.py index 35154a4b..22f722fc 100644 --- a/libcloudforensics/providers/gcp/internal/compute.py +++ b/libcloudforensics/providers/gcp/internal/compute.py @@ -525,7 +525,9 @@ def CreateDiskFromSnapshot( snapshot: 'GoogleComputeSnapshot', disk_name: Optional[str] = None, disk_name_prefix: Optional[str] = None, - disk_type: str = 'pd-standard') -> 'GoogleComputeDisk': + disk_type: str = 'pd-standard', + dest_project: Optional[str] = None, + dest_zone: Optional[str] = None) -> 'GoogleComputeDisk': """Create a new disk based on a Snapshot. Args: @@ -536,6 +538,10 @@ def CreateDiskFromSnapshot( which disk type to use to create the disk. Default is pd-standard. Use pd-ssd to have a SSD disk. You can list all available disk types by running the following command: gcloud compute disk-types list + dest_project (str): Optional. The destination project where the disk + should be created. Default is self.project_id. + dest_zone (str): Optional. The destination zone where the disk + should be created. Default is self.default_zone. Returns: GoogleComputeDisk: Google Compute Disk. @@ -547,6 +553,10 @@ def CreateDiskFromSnapshot( if not disk_name: disk_name = common.GenerateDiskName(snapshot, disk_name_prefix) + + project_id = dest_project or self.project_id + zone = dest_zone or self.default_zone + body = { 'name': disk_name, @@ -554,12 +564,12 @@ def CreateDiskFromSnapshot( snapshot.GetSourceString(), 'type': 'projects/{0:s}/zones/{1:s}/diskTypes/{2:s}'.format( - self.project_id, self.default_zone, disk_type) + project_id, zone, disk_type) } try: gce_disks_client = self.GceApi().disks() # pylint: disable=no-member request = gce_disks_client.insert( - project=self.project_id, zone=self.default_zone, body=body) + project=project_id, zone=zone, body=body) response = request.execute() except HttpError as exception: if exception.resp.status == 409: @@ -570,9 +580,9 @@ def CreateDiskFromSnapshot( 'Unknown error occurred when creating disk from Snapshot:' ' {0!s}'.format(exception), __name__) from exception - self.BlockOperation(response, zone=self.default_zone) + self.BlockOperation(response, zone=zone) return GoogleComputeDisk( - project_id=self.project_id, zone=self.default_zone, name=disk_name) + project_id=project_id, zone=zone, name=disk_name) def GetMachineTypes(self, machine_type: str, zone: Optional[str] = None) -> Dict[str, Any]: diff --git a/tests/providers/gcp/internal/test_compute.py b/tests/providers/gcp/internal/test_compute.py index 5de80e0b..f4546730 100644 --- a/tests/providers/gcp/internal/test_compute.py +++ b/tests/providers/gcp/internal/test_compute.py @@ -173,6 +173,15 @@ def testCreateDiskFromSnapshot(self, mock_gce_api, mock_block_operation): disk_from_snapshot, compute.GoogleComputeDisk) self.assertEqual('new-forensics-disk', disk_from_snapshot.name) + # CreateDiskFromSnapshot(Snapshot=gcp_mocks.FAKE_SNAPSHOT, + # dest_project='other-project', dest_zone='other-zone') + disk_from_snapshot = gcp_mocks.FAKE_ANALYSIS_PROJECT.compute.CreateDiskFromSnapshot( + gcp_mocks.FAKE_SNAPSHOT, dest_project='other-project', dest_zone='other-zone') + self.assertIsInstance( + disk_from_snapshot, compute.GoogleComputeDisk) + self.assertEqual('other-project', disk_from_snapshot.project_id) + self.assertEqual('other-zone', disk_from_snapshot.zone) + # CreateDiskFromSnapshot(Snapshot=gcp_mocks.FAKE_SNAPSHOT, # disk_name='fake-disk') where 'fake-disk' exists already disks.return_value.insert.return_value.execute.side_effect = HttpError(