Add VM shutdown, capture template, create from template

vmdriver.py

@@ -26,6 +26,7 @@ SECRET = os.getenv('SECRET')
 TENANT = os.getenv('TENANT')
 GROUP_NAME = os.getenv('GROUP_NAME')
 REGION = os.getenv('REGION')
+STORAGE_NAME = os.getenv('STORAGE_NAME')
 
 credentials = ServicePrincipalCredentials(
     client_id = CLIENT_ID,
@@ -192,6 +193,26 @@ def create(vm_desc):
             )
         )
 
+    vhd_uri = None
+    image_reference = None
+    os_disk_type = None
+
+    if vm_desc["vhd_uri"]:
+        #TODO: select OS Disk type based on OS
+        os_disk_type = azure.mgmt.compute.models.OperatingSystemTypes.linux
+        vhd_uri = azure.mgmt.compute.models.VirtualHardDisk(
+            vm_desc["vhd_uri"]
+        )
+        logging.info("creating vm from template")
+    else:
+        image_reference = azure.mgmt.compute.models.ImageReference(
+            publisher=vm_desc["os_publisher"],
+            offer=vm_desc["os_offer"],
+            sku=vm_desc["os_sku"],
+            version=vm_desc["os_version"],
+        )
+        logging.info("creating vm from Azure image")
+
     poller = compute_client.virtual_machines.create_or_update(
         GROUP_NAME,
         vm_desc["name"],
@@ -210,6 +231,7 @@ def create(vm_desc):
             ),
             storage_profile=azure.mgmt.compute.models.StorageProfile(
                 os_disk=azure.mgmt.compute.models.OSDisk(
+                    os_type=os_disk_type,
                     caching=azure.mgmt.compute.models.CachingTypes.none,
                     create_option=\
                     azure.mgmt.compute.models.DiskCreateOptionTypes.from_image,
@@ -220,13 +242,9 @@ def create(vm_desc):
                             os_disk_name,
                         ),
                     ),
+                    image = vhd_uri,
                 ),
-                image_reference = azure.mgmt.compute.models.ImageReference(
-                    publisher=vm_desc["os_publisher"],
-                    offer=vm_desc["os_offer"],
-                    sku=vm_desc["os_sku"],
-                    version=vm_desc["os_version"],
-                ),
+                image_reference = image_reference,
             ),
         ),
     )
@@ -239,42 +257,46 @@ def create(vm_desc):
     except Exception, e:
         logging.error("cloud not create vm '%s'" % vm_desc["name"])
         logging.error(str(e))
+        raise e
+
+@celery.task
+def shutdown(vm_name):
+    """ Shutdown a virtual machine in azure. """
+    poller = compute_client.virtual_machines.deallocate(GROUP_NAME, vm_name)
+    try:
+        poller.wait()
+        logging.info("shutdown vm: %s" % vm_name)
+        return "ok"
+    except Exception, e:
+        logging.error("cloud not stop vm '%s'" % vm_name)
         return None
 
+@celery.task
+def save_as_template(vm_name, vm_pk):
+    """ Save virtual machine as template in azure. """
+    try:
+        compute_client.virtual_machines.generalize(
+            GROUP_NAME,
+            vm_name
+        )
 
-class shutdown(AbortableTask):
-    """ Shutdown virtual machine (need ACPI support).
-    Return When domain is missiing.
-    This job is abortable:
-        AbortableAsyncResult(id="<<jobid>>").abort()
-    """
-    time_limit = 120
+        os_disk_name = "%s_os_disk_template" % vm_pk
 
-    @req_connection
-    def run(self, args):
-        from time import sleep
-        name, = args
-        try:
-            domain = lookupByName(name)
-            domain.shutdown()
-            while True:
-                try:
-                    Connection.get().lookupByName(name)
-                except libvirt.libvirtError as e:
-                    if e.get_error_code() == libvirt.VIR_ERR_NO_DOMAIN:
-                        return
-                    else:
-                        raise
-                else:
-                    if self.is_aborted():
-                        logging.info("Shutdown aborted on vm: %s", name)
-                        return
-                    sleep(5)
-        except libvirt.libvirtError as e:
-            new_e = Exception(e.get_error_message())
-            new_e.libvirtError = True
-            raise new_e
+        result = compute_client.virtual_machines.capture(
+            GROUP_NAME,
+            vm_name,
+            azure.mgmt.compute.models.VirtualMachineCaptureParameters(
+                os_disk_name,
+                "templates",
+                False
+            )
+        )
+        result.wait()
 
+        return result.result().output["resources"][0]["properties"]\
+        ["storageProfile"]["osDisk"]["image"]["uri"]
+    except Exception, e:
+        raise
 
 @celery.task
 @req_connection