Additional provider features added and fixed some bugs.

* Added support for SSH keys, image passwords, SSD disk type, and CPU family.
* Adjusted server create so that IP address is returned in response.
* Restructured remove server method(s) to handle change status properly, gracefully handle missing servers, and improve overall performance.
* Prevent duplicate server names from being provisioned so removals can be handled appropriately.
* Fixed a bug in the count increment being a string rather than an integer.
* Fixed issue with create_volume returning invalid response.
* Fixed type bug in volume instance_ids for volume removal and improved volume management.
* Fixed type bug in instance_ids for proper server removal and moved boot volume creation into composite server build request.
* General clean up.

2 files changed, 222 insertions, 155 deletions

diff --git a/cloud/profitbricks/profitbricks.py b/cloud/profitbricks/profitbricks.py
index 556c6528..0a1e2999 100644
--- a/cloud/profitbricks/profitbricks.py
+++ b/cloud/profitbricks/profitbricks.py
@@ -31,13 +31,21 @@ options:
     description:
       - The name of the virtual machine.
     required: true
-  image: 
+  image:
     description:
       - The system image ID for creating the virtual machine, e.g. a3eae284-a2fe-11e4-b187-5f1f641608c8.
     required: true
+  image_password:
+    description:
+      - Password set for the administrative user.
+    required: false
+  ssh_keys:
+    description:
+      - Public SSH keys allowing access to the virtual machine.
+    required: false
   datacenter:
     description:
-      - The Datacenter to provision this virtual machine.
+      - The datacenter to provision this virtual machine.
     required: false
     default: null
   cores:
@@ -50,6 +58,12 @@ options:
       - The amount of memory to allocate to the virtual machine.
     required: false
     default: 2048
+  cpu_family:
+    description:
+      - The CPU family type to allocate to the virtual machine.
+    required: false
+    default: AMD_OPTERON
+    choices: [ "AMD_OPTERON", "INTEL_XEON" ]
   volume_size:
     description:
       - The size in GB of the boot volume.
@@ -72,7 +86,7 @@ options:
     default: 1
   location:
     description:
-      - The datacenter location. Use only if you want to create the Datacenter or else this value is ignored. 
+      - The datacenter location. Use only if you want to create the Datacenter or else this value is ignored.
     required: false
     default: us/las
     choices: [ "us/las", "us/lasdev", "de/fra", "de/fkb" ]
@@ -129,7 +143,7 @@ EXAMPLES = '''
 
 # Note: These examples do not set authentication details, see the AWS Guide for details.
 
-# Provisioning example. This will create three servers and enumerate their names. 
+# Provisioning example. This will create three servers and enumerate their names.
 
 - profitbricks:
     datacenter: Tardis One
@@ -137,6 +151,7 @@ EXAMPLES = '''
     cores: 4
     ram: 2048
     volume_size: 50
+    cpu_family: INTEL_XEON
     image: a3eae284-a2fe-11e4-b187-5f1f641608c8
     location: us/las
     count: 3
@@ -218,11 +233,15 @@ def _wait_for_completion(profitbricks, promise, wait_timeout, msg):
             promise['requestId']
             ) + '" to complete.')
 
+
 def _create_machine(module, profitbricks, datacenter, name):
-    image = module.params.get('image')
     cores = module.params.get('cores')
     ram = module.params.get('ram')
+    cpu_family = module.params.get('cpu_family')
     volume_size = module.params.get('volume_size')
+    disk_type = module.params.get('disk_type')
+    image_password = module.params.get('image_password')
+    ssh_keys = module.params.get('ssh_keys')
     bus = module.params.get('bus')
     lan = module.params.get('lan')
     assign_public_ip = module.params.get('assign_public_ip')
@@ -234,26 +253,6 @@ def _create_machine(module, profitbricks, datacenter, name):
     wait = module.params.get('wait')
     wait_timeout = module.params.get('wait_timeout')
 
-    try:
-        # Generate name, but grab first 10 chars so we don't
-        # screw up the uuid match routine.
-        v = Volume(
-            name=str(uuid.uuid4()).replace('-','')[:10],
-            size=volume_size,
-            image=image,
-            bus=bus)
-
-        volume_response = profitbricks.create_volume(
-            datacenter_id=datacenter, volume=v)
-
-        # We're forced to wait on the volume creation since
-        # server create relies upon this existing.
-
-        _wait_for_completion(profitbricks, volume_response,
-                             wait_timeout, "create_volume")
-    except Exception as e:
-        module.fail_json(msg="failed to create the new volume: %s" % str(e))
-
     if assign_public_ip:
         public_found = False
 
@@ -269,81 +268,64 @@ def _create_machine(module, profitbricks, datacenter, name):
                 public=True)
 
             lan_response = profitbricks.create_lan(datacenter, i)
-
-            lan = lan_response['id']
-
             _wait_for_completion(profitbricks, lan_response,
                                  wait_timeout, "_create_machine")
+            lan = lan_response['id']
 
-    try:
-        n = NIC(
-            lan=int(lan)
-            )
-
-        nics = [n]
+    v = Volume(
+        name=str(uuid.uuid4()).replace('-', '')[:10],
+        size=volume_size,
+        image=image,
+        image_password=image_password,
+        ssh_keys=ssh_keys,
+        disk_type=disk_type,
+        bus=bus)
+
+    n = NIC(
+        lan=int(lan)
+        )
 
-        s = Server(
-            name=name,
-            ram=ram,
-            cores=cores,
-            nics=nics,
-            boot_volume_id=volume_response['id']
-            )
+    s = Server(
+        name=name,
+        ram=ram,
+        cores=cores,
+        cpu_family=cpu_family,
+        create_volumes=[v],
+        nics=[n],
+        )
 
-        server_response = profitbricks.create_server(
+    try:
+        create_server_response = profitbricks.create_server(
             datacenter_id=datacenter, server=s)
 
-        if wait:
-            _wait_for_completion(profitbricks, server_response,
-                                 wait_timeout, "create_virtual_machine")
-
+        _wait_for_completion(profitbricks, create_server_response,
+                             wait_timeout, "create_virtual_machine")
 
-        return (server_response)
+        server_response = profitbricks.get_server(
+            datacenter_id=datacenter,
+            server_id=create_server_response['id'],
+            depth=3
+        )
     except Exception as e:
         module.fail_json(msg="failed to create the new server: %s" % str(e))
+    else:
+        return server_response
 
-def _remove_machine(module, profitbricks, datacenter, name):
-    remove_boot_volume = module.params.get('remove_boot_volume')
-    wait = module.params.get('wait')
-    wait_timeout = module.params.get('wait_timeout')
-    changed = False
-
-    # User provided the actual UUID instead of the name.
-    try:
-        if remove_boot_volume:
-            # Collect information needed for later. 
-            server = profitbricks.get_server(datacenter, name)
-            volume_id = server['properties']['bootVolume']['href'].split('/')[7]
-
-        server_response = profitbricks.delete_server(datacenter, name)
-        changed = True
-
-    except Exception as e:
-        module.fail_json(msg="failed to terminate the virtual server: %s" % str(e))
-
-    # Remove the bootVolume
-    if remove_boot_volume:
-        try:
-            volume_response = profitbricks.delete_volume(datacenter, volume_id)
-
-        except Exception as e:
-            module.fail_json(msg="failed to remove the virtual server's bootvolume: %s" % str(e))
-
-    return changed
 
-def _startstop_machine(module, profitbricks, datacenter, name):
+def _startstop_machine(module, profitbricks, datacenter_id, server_id):
     state = module.params.get('state')
 
     try:
         if state == 'running':
-            profitbricks.start_server(datacenter, name)
+            profitbricks.start_server(datacenter_id, server_id)
         else:
-            profitbricks.stop_server(datacenter, name)
+            profitbricks.stop_server(datacenter_id, server_id)
 
         return True
     except Exception as e:
         module.fail_json(msg="failed to start or stop the virtual machine %s: %s" % (name, str(e)))
 
+
 def _create_datacenter(module, profitbricks):
     datacenter = module.params.get('datacenter')
     location = module.params.get('location')
@@ -364,6 +346,7 @@ def _create_datacenter(module, profitbricks):
     except Exception as e:
         module.fail_json(msg="failed to create the new server(s): %s" % str(e))
 
+
 def create_virtual_machine(module, profitbricks):
     """
     Create new virtual machine
@@ -386,19 +369,15 @@ def create_virtual_machine(module, profitbricks):
     virtual_machines = []
     virtual_machine_ids = []
 
-    # Locate UUID for Datacenter
-    if not (uuid_match.match(datacenter)):
-        datacenter_list = profitbricks.list_datacenters()
-        for d in datacenter_list['items']:
-            dc = profitbricks.get_datacenter(d['id'])
-            if datacenter == dc['properties']['name']:
-                datacenter = d['id']
-                datacenter_found = True
-                break
+    # Locate UUID for datacenter if referenced by name.
+    datacenter_list = profitbricks.list_datacenters()
+    datacenter_id = _get_datacenter_id(datacenter_list, datacenter)
+    if datacenter_id:
+        datacenter_found = True
 
     if not datacenter_found:
         datacenter_response = _create_datacenter(module, profitbricks)
-        datacenter = datacenter_response['id']
+        datacenter_id = datacenter_response['id']
 
         _wait_for_completion(profitbricks, datacenter_response,
                              wait_timeout, "create_virtual_machine")
@@ -415,24 +394,31 @@ def create_virtual_machine(module, profitbricks):
             else:
                 module.fail_json(msg=e.message)
 
-        number_range = xrange(count_offset,count_offset + count + len(numbers))
+        number_range = xrange(count_offset, count_offset + count + len(numbers))
         available_numbers = list(set(number_range).difference(numbers))
         names = []
         numbers_to_use = available_numbers[:count]
         for number in numbers_to_use:
             names.append(name % number)
     else:
-        names = [name] * count
+        names = [name]
 
-    for name in  names: 
-        create_response = _create_machine(module, profitbricks, str(datacenter), name)
-        nics = profitbricks.list_nics(datacenter,create_response['id'])
+    # Prefetch a list of servers for later comparison.
+    server_list = profitbricks.list_servers(datacenter_id)
+    for name in names:
+        # Skip server creation if the server already exists.
+        if _get_server_id(server_list, name):
+            continue
+
+        create_response = _create_machine(module, profitbricks, str(datacenter_id), name)
+        nics = profitbricks.list_nics(datacenter_id, create_response['id'])
         for n in nics['items']:
             if lan == n['properties']['lan']:
-                create_response.update({ 'public_ip': n['properties']['ips'][0] })
+                create_response.update({'public_ip': n['properties']['ips'][0]})
 
         virtual_machines.append(create_response)
-        failed = False
+
+    failed = False
 
     results = {
         'failed': failed,
@@ -445,9 +431,10 @@ def create_virtual_machine(module, profitbricks):
 
     return results
 
+
 def remove_virtual_machine(module, profitbricks):
     """
-    Removes a virtual machine. 
+    Removes a virtual machine.
 
     This will remove the virtual machine along with the bootVolume.
 
@@ -459,36 +446,56 @@ def remove_virtual_machine(module, profitbricks):
     Returns:
         True if a new virtual server was deleted, false otherwise
     """
+    datacenter = module.params.get('datacenter')
+    instance_ids = module.params.get('instance_ids')
+    remove_boot_volume = module.params.get('remove_boot_volume')
+    changed = False
+
     if not isinstance(module.params.get('instance_ids'), list) or len(module.params.get('instance_ids')) < 1:
         module.fail_json(msg='instance_ids should be a list of virtual machine ids or names, aborting')
 
-    datacenter = module.params.get('datacenter')
-    instance_ids = module.params.get('instance_ids')
+    # Locate UUID for datacenter if referenced by name.
+    datacenter_list = profitbricks.list_datacenters()
+    datacenter_id = _get_datacenter_id(datacenter_list, datacenter)
+    if not datacenter_id:
+        module.fail_json(msg='Virtual data center \'%s\' not found.' % str(datacenter))
+
+    # Prefetch server list for later comparison.
+    server_list = profitbricks.list_servers(datacenter_id)
+    for instance in instance_ids:
+        # Locate UUID for server if referenced by name.
+        server_id = _get_server_id(server_list, instance)
+        if server_id:
+            # Remove the server's boot volume
+            if remove_boot_volume:
+                _remove_boot_volume(module, profitbricks, datacenter_id, server_id)
+
+            # Remove the server
+            try:
+                server_response = profitbricks.delete_server(datacenter_id, server_id)
+            except Exception as e:
+                module.fail_json(msg="failed to terminate the virtual server: %s" % str(e))
+            else:
+                changed = True
 
-    # Locate UUID for Datacenter
-    if not (uuid_match.match(datacenter)):
-        datacenter_list = profitbricks.list_datacenters()
-        for d in datacenter_list['items']:
-            dc = profitbricks.get_datacenter(d['id'])
-            if datacenter == dc['properties']['name']:
-                datacenter = d['id']
-                break
+    return changed
 
-    for n in instance_ids:
-        if(uuid_match.match(n)):
-            _remove_machine(module, profitbricks, d['id'], n)
-        else:
-            servers = profitbricks.list_servers(d['id'])
 
-            for s in servers['items']:
-                if n == s['properties']['name']:
-                    server_id = s['id']
+def _remove_boot_volume(module, profitbricks, datacenter_id, server_id):
+    """
+    Remove the boot volume from the server
+    """
+    try:
+        server = profitbricks.get_server(datacenter_id, server_id)
+        volume_id = server['properties']['bootVolume']['id']
+        volume_response = profitbricks.delete_volume(datacenter_id, volume_id)
+    except Exception as e:
+        module.fail_json(msg="failed to remove the server's boot volume: %s" % str(e))
 
-                    _remove_machine(module, profitbricks, datacenter, server_id)
 
 def startstop_machine(module, profitbricks, state):
     """
-    Starts or Stops a virtual machine. 
+    Starts or Stops a virtual machine.
 
     module : AnsibleModule object
     profitbricks: authenticated profitbricks object.
@@ -506,41 +513,32 @@ def startstop_machine(module, profitbricks, state):
     datacenter = module.params.get('datacenter')
     instance_ids = module.params.get('instance_ids')
 
-    # Locate UUID for Datacenter
-    if not (uuid_match.match(datacenter)):
-        datacenter_list = profitbricks.list_datacenters()
-        for d in datacenter_list['items']:
-            dc = profitbricks.get_datacenter(d['id'])
-            if datacenter == dc['properties']['name']:
-                datacenter = d['id']
-                break
-
-    for n in instance_ids:
-        if(uuid_match.match(n)):
-            _startstop_machine(module, profitbricks, datacenter, n)
-
+    # Locate UUID for datacenter if referenced by name.
+    datacenter_list = profitbricks.list_datacenters()
+    datacenter_id = _get_datacenter_id(datacenter_list, datacenter)
+    if not datacenter_id:
+        module.fail_json(msg='Virtual data center \'%s\' not found.' % str(datacenter))
+
+    # Prefetch server list for later comparison.
+    server_list = profitbricks.list_servers(datacenter_id)
+    for instance in instance_ids:
+        # Locate UUID of server if referenced by name.
+        server_id = _get_server_id(server_list, instance)
+        if server_id:
+            _startstop_machine(module, profitbricks, datacenter_id, server_id)
             changed = True
-        else:
-            servers = profitbricks.list_servers(d['id'])
-
-            for s in servers['items']:
-                if n == s['properties']['name']:
-                    server_id = s['id']
-                    _startstop_machine(module, profitbricks, datacenter, server_id)
-
-                    changed = True
 
     if wait:
         wait_timeout = time.time() + wait_timeout
         while wait_timeout > time.time():
             matched_instances = []
-            for res in profitbricks.list_servers(datacenter)['items']:
+            for res in profitbricks.list_servers(datacenter_id)['items']:
                 if state == 'running':
                     if res['properties']['vmState'].lower() == state:
                         matched_instances.append(res)
                 elif state == 'stopped':
                     if res['properties']['vmState'].lower() == 'shutoff':
-                        matched_instances.append(res)                    
+                        matched_instances.append(res)
 
             if len(matched_instances) < len(instance_ids):
                 time.sleep(5)
@@ -549,10 +547,31 @@ def startstop_machine(module, profitbricks, state):
 
         if wait_timeout <= time.time():
             # waiting took too long
-            module.fail_json(msg = "wait for virtual machine state timeout on %s" % time.asctime())
+            module.fail_json(msg="wait for virtual machine state timeout on %s" % time.asctime())
 
     return (changed)
 
+
+def _get_datacenter_id(datacenters, identity):
+    """
+    Fetch and return datacenter UUID by datacenter name if found.
+    """
+    for datacenter in datacenters['items']:
+        if identity in (datacenter['properties']['name'], datacenter['id']):
+            return datacenter['id']
+    return None
+
+
+def _get_server_id(servers, identity):
+    """
+    Fetch and return server UUID by server name if found.
+    """
+    for server in servers['items']:
+        if identity in (server['properties']['name'], server['id']):
+            return server['id']
+    return None
+
+
 def main():
     module = AnsibleModule(
         argument_spec=dict(
@@ -561,12 +580,16 @@ def main():
             image=dict(),
             cores=dict(default=2),
             ram=dict(default=2048),
+            cpu_family=dict(default='AMD_OPTERON'),
             volume_size=dict(default=10),
+            disk_type=dict(default='HDD'),
+            image_password=dict(default=None),
+            ssh_keys=dict(type='list', default=[]),
             bus=dict(default='VIRTIO'),
             lan=dict(default=1),
-            count=dict(default=1),
+            count=dict(type='int', default=1),
             auto_increment=dict(type='bool', default=True),
-            instance_ids=dict(),
+            instance_ids=dict(type='list', default=[]),
             subscription_user=dict(),
             subscription_password=dict(),
             location=dict(choices=LOCATIONS, default='us/las'),
@@ -594,7 +617,7 @@ def main():
 
     if state == 'absent':
         if not module.params.get('datacenter'):
-            module.fail_json(msg='datacenter parameter is required ' + 
+            module.fail_json(msg='datacenter parameter is required ' +
                 'for running or stopping machines.')
 
         try:
@@ -605,7 +628,7 @@ def main():
 
     elif state in ('running', 'stopped'):
         if not module.params.get('datacenter'):
-            module.fail_json(msg='datacenter parameter is required for ' + 
+            module.fail_json(msg='datacenter parameter is required for ' +
                 'running or stopping machines.')
         try:
             (changed) = startstop_machine(module, profitbricks, state)
@@ -619,10 +642,10 @@ def main():
         if not module.params.get('image'):
             module.fail_json(msg='image parameter is required for new instance')
         if not module.params.get('subscription_user'):
-            module.fail_json(msg='subscription_user parameter is ' + 
+            module.fail_json(msg='subscription_user parameter is ' +
                 'required for new instance')
         if not module.params.get('subscription_password'):
-            module.fail_json(msg='subscription_password parameter is ' + 
+            module.fail_json(msg='subscription_password parameter is ' +
                 'required for new instance')
 
         try:
@@ -634,4 +657,3 @@ def main():
 from ansible.module_utils.basic import *
 
 main()
-
diff --git a/cloud/profitbricks/profitbricks_volume.py b/cloud/profitbricks/profitbricks_volume.py
index 89a69d5e..802511cc 100644
--- a/cloud/profitbricks/profitbricks_volume.py
+++ b/cloud/profitbricks/profitbricks_volume.py
@@ -45,11 +45,20 @@ options:
     description:
       - The system image ID for the volume, e.g. a3eae284-a2fe-11e4-b187-5f1f641608c8. This can also be a snapshot image ID.
     required: true
+  image_password:
+    description:
+      - Password set for the administrative user.
+    required: false
+  ssh_keys:
+    description:
+      - Public SSH keys allowing access to the virtual machine.
+    required: false
   disk_type:
     description:
-      - The disk type. Currently only HDD.
+      - The disk type of the volume.
     required: false
     default: HDD
+    choices: [ "HDD", "SSD" ]
   licence_type:
     description:
       - The licence type for the volume. This is used when the image is non-standard. 
@@ -163,6 +172,8 @@ def _create_volume(module, profitbricks, datacenter, name):
     size = module.params.get('size')
     bus = module.params.get('bus')
     image = module.params.get('image')
+    image_password = module.params.get('image_password')
+    ssh_keys = module.params.get('ssh_keys')
     disk_type = module.params.get('disk_type')
     licence_type = module.params.get('licence_type')
     wait_timeout = module.params.get('wait_timeout')
@@ -174,6 +185,8 @@ def _create_volume(module, profitbricks, datacenter, name):
             size=size,
             bus=bus,
             image=image,
+            image_password=image_password,
+            ssh_keys=ssh_keys,
             disk_type=disk_type,
             licence_type=licence_type
             )
@@ -250,9 +263,10 @@ def create_volume(module, profitbricks):
     else:
         names = [name] * count
 
-    for name in  names: 
+    for name in names: 
         create_response = _create_volume(module, profitbricks, str(datacenter), name)
         volumes.append(create_response)
+        _attach_volume(module, profitbricks, datacenter, create_response['id'])
         failed = False
 
     results = {
@@ -308,19 +322,50 @@ def delete_volume(module, profitbricks):
 
     return changed
 
+def _attach_volume(module, profitbricks, datacenter, volume):
+    """
+    Attaches a volume.
+
+    This will attach a volume to the server.
+
+    module : AnsibleModule object
+    profitbricks: authenticated profitbricks object.
+
+    Returns:
+        True if the volume was attached, false otherwise
+    """
+    server = module.params.get('server')
+
+    # Locate UUID for Server
+    if server: 
+        if not (uuid_match.match(server)):
+            server_list = profitbricks.list_servers(datacenter)
+            for s in server_list['items']:
+                if server == s['properties']['name']:
+                    server= s['id']
+                    break
+
+        try:
+            return profitbricks.attach_volume(datacenter, server, volume)
+        except Exception as e:
+            module.fail_json(msg='failed to attach volume: %s' % str(e))
+
 def main():
     module = AnsibleModule(
         argument_spec=dict(
             datacenter=dict(),
+            server=dict(),
             name=dict(),
             size=dict(default=10),
             bus=dict(default='VIRTIO'),
             image=dict(),
+            image_password=dict(default=None),
+            ssh_keys=dict(type='list', default=[]),
             disk_type=dict(default='HDD'),
             licence_type=dict(default='UNKNOWN'),
-            count=dict(default=1),
+            count=dict(type='int', default=1),
             auto_increment=dict(type='bool', default=True),
-            instance_ids=dict(),
+            instance_ids=dict(type='list', default=[]),
             subscription_user=dict(),
             subscription_password=dict(),
             wait=dict(type='bool', default=True),
@@ -360,11 +405,11 @@ def main():
             module.fail_json(msg='name parameter is required for new instance')
 
         try:
-            (failed, volume_dict_array) = create_volume(module, profitbricks)
-            module.exit_json(failed=failed, volumes=volume_dict_array)
+            (volume_dict_array) = create_volume(module, profitbricks)
+            module.exit_json(**volume_dict_array)
         except Exception as e:
             module.fail_json(msg='failed to set volume state: %s' % str(e))
 
 from ansible.module_utils.basic import *
 
-main()
\ No newline at end of file
+main()