Add windows support to AWS

perfkitbenchmarker/aws/aws_virtual_machine.py

@@ -18,6 +18,7 @@
 operate on the VM: boot, shutdown, etc.
 """
 
+import base64
 import json
 import logging
 import threading
@@ -28,6 +29,7 @@
 from perfkitbenchmarker import linux_virtual_machine
 from perfkitbenchmarker import virtual_machine
 from perfkitbenchmarker import vm_util
+from perfkitbenchmarker import windows_virtual_machine
 from perfkitbenchmarker.aws import aws_disk
 from perfkitbenchmarker.aws import aws_network
 from perfkitbenchmarker.aws import util
@@ -72,6 +74,7 @@
         SA_EAST_1: 'ami-6770d87a',
     }
 }
+HVM_US_EAST_1_WINDOWS_AMI = 'ami-cc93a8a4'
 PLACEMENT_GROUP_PREFIXES = frozenset(
     ['c3', 'c4', 'cc2', 'cg1', 'g2', 'cr1', 'r3', 'hi1', 'i2'])
 NUM_LOCAL_VOLUMES = {
@@ -142,6 +145,7 @@ def __init__(self, vm_spec):
     self.network = aws_network.AwsNetwork.GetNetwork(self.zone)
     if self.machine_type in NUM_LOCAL_VOLUMES:
       self.max_local_disks = NUM_LOCAL_VOLUMES[self.machine_type]
+    self.user_data = None
 
   @classmethod
   def SetVmSpecDefaults(cls, vm_spec):
@@ -239,6 +243,8 @@ def _Create(self):
         '--key-name=%s' % 'perfkit-key-%s' % FLAGS.run_uri]
     if block_device_map:
       create_cmd.append('--block-device-mappings=%s' % block_device_map)
+    if self.user_data:
+      create_cmd.append('--user-data=%s' % self.user_data)
     stdout, _, _ = vm_util.IssueCommand(create_cmd)
     response = json.loads(stdout)
     self.id = response['Instances'][0]['InstanceId']
@@ -278,19 +284,24 @@ def CreateScratchDisk(self, disk_spec):
       disk_spec: virtual_machine.BaseDiskSpec object of the disk.
     """
     # Instantiate the disk(s) that we want to create.
-    if disk_spec.disk_type == disk.LOCAL:
-      disks = []
-      for _ in range(disk_spec.num_striped_disks):
-        local_disk = aws_disk.AwsDisk(disk_spec, self.zone)
-        local_disk.device_letter = chr(ord(DRIVE_START_LETTER) +
-                                       self.local_disk_counter)
+    disks = []
+    for _ in range(disk_spec.num_striped_disks):
+      data_disk = aws_disk.AwsDisk(disk_spec, self.zone)
+      if disk_spec.disk_type == disk.LOCAL:
+        data_disk.device_letter = chr(ord(DRIVE_START_LETTER) +
+                                      self.local_disk_counter)
+        # Local disk numbers start at 1 (0 is the system disk).
+        data_disk.disk_number = self.local_disk_counter + 1
         self.local_disk_counter += 1
-        disks.append(local_disk)
-      if self.local_disk_counter > self.max_local_disks:
-        raise errors.Error('Not enough local disks.')
-    else:
-      disks = [aws_disk.AwsDisk(disk_spec, self.zone)
-               for _ in range(disk_spec.num_striped_disks)]
+        if self.local_disk_counter > self.max_local_disks:
+          raise errors.Error('Not enough local disks.')
+      else:
+        # Remote disk numbers start at 1 + max_local disks (0 is the system disk
+        # and local disks occupy [1, max_local_disks]).
+        data_disk.disk_number = (self.remote_disk_counter +
+                                 1 + self.max_local_disks)
+        self.remote_disk_counter += 1
+      disks.append(data_disk)
 
     self._CreateScratchDiskFromDisks(disk_spec, disks)
 
@@ -328,3 +339,67 @@ def _GetDefaultImage(machine_type, region):
 class RhelBasedAwsVirtualMachine(AwsVirtualMachine,
                                  linux_virtual_machine.RhelMixin):
   pass
+
+
+class WindowsAwsVirtualMachine(AwsVirtualMachine,
+                               windows_virtual_machine.WindowsMixin):
+
+  def __init__(self, vm_spec):
+    super(WindowsAwsVirtualMachine, self).__init__(vm_spec)
+    self.user_name = 'Administrator'
+    self.user_data = ('<powershell>%s</powershell>' %
+                      windows_virtual_machine.STARTUP_SCRIPT)
+
+  @staticmethod
+  def _GetDefaultImage(machine_type, region):
+    """Returns the default image given the machine type and region.
+
+    If no default is configured, this will return None.
+    """
+    prefix = machine_type.split('.')[0]
+    if prefix in NON_HVM_PREFIXES or region != US_EAST_1:
+      return None
+    return HVM_US_EAST_1_WINDOWS_AMI
+
+  @vm_util.Retry()
+  def _GetDecodedPasswordData(self):
+    # Retreive a base64 encoded, encrypted password for the VM.
+    get_password_cmd = util.AWS_PREFIX + [
+        'ec2',
+        'get-password-data',
+        '--region=%s' % self.region,
+        '--instance-id=%s' % self.id]
+    stdout, _ = util.IssueRetryableCommand(get_password_cmd)
+    response = json.loads(stdout)
+    password_data = response['PasswordData']
+
+    # AWS may not populate the password data until some time after
+    # the VM shows as running. Simply retry until the data shows up.
+    if not password_data:
+      raise ValueError('No PasswordData in response.')
+
+    # Decode the password data.
+    return base64.b64decode(password_data)
+
+
+  def _PostCreate(self):
+    """Retrieve generic VM info and then retrieve the VM's password."""
+    super(WindowsAwsVirtualMachine, self)._PostCreate()
+
+    # Get the decoded password data.
+    decoded_password_data = self._GetDecodedPasswordData()
+
+    # Write the encrypted data to a file, and use openssl to
+    # decrypt the password.
+    with vm_util.NamedTemporaryFile() as tf:
+      tf.write(decoded_password_data)
+      tf.close()
+      decrypt_cmd = ['openssl',
+                     'rsautl',
+                     '-decrypt',
+                     '-in',
+                     tf.name,
+                     '-inkey',
+                     vm_util.GetPrivateKeyPath()]
+      password, _ = vm_util.IssueRetryableCommand(decrypt_cmd)
+      self.password = password

perfkitbenchmarker/benchmark_spec.py

@@ -68,7 +68,8 @@
     AWS: {
         VIRTUAL_MACHINE: {
             DEBIAN: aws_virtual_machine.DebianBasedAwsVirtualMachine,
-            RHEL: aws_virtual_machine.RhelBasedAwsVirtualMachine
+            RHEL: aws_virtual_machine.RhelBasedAwsVirtualMachine,
+            WINDOWS: aws_virtual_machine.WindowsAwsVirtualMachine
         },
         FIREWALL: aws_network.AwsFirewall
     },