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
