From 7bca6a56c9ca7e0632f58e08c73140cb359bebd1 Mon Sep 17 00:00:00 2001
From: awkspace <self@awk.space>
Date: Sun, 24 Feb 2019 13:57:19 -0500
Subject: [PATCH] Add snapshot support for vultr_server (#52781)

---
 .../modules/cloud/vultr/vultr_server.py       | 38 ++++++++++++++++---
 1 file changed, 33 insertions(+), 5 deletions(-)

diff --git a/lib/ansible/modules/cloud/vultr/vultr_server.py b/lib/ansible/modules/cloud/vultr/vultr_server.py
index 024cbff6b20..9f6cc7a267d 100644
--- a/lib/ansible/modules/cloud/vultr/vultr_server.py
+++ b/lib/ansible/modules/cloud/vultr/vultr_server.py
@@ -32,7 +32,11 @@ options:
   os:
     description:
       - The operating system.
-      - Required if the server does not yet exist.
+      - Required if the server does not yet exist and is not restoring from a snapshot.
+  snapshot:
+    version_added: "2.8"
+    description:
+      - Name of snapshot to restore server from.
   firewall_group:
     description:
       - The firewall group to assign this server to.
@@ -379,13 +383,26 @@ class AnsibleVultrServer(Vultr):
         )
 
     def get_os(self):
+        if self.module.params.get('snapshot'):
+            os_name = 'Snapshot'
+        else:
+            os_name = self.module.params.get('os')
+
         return self.query_resource_by_key(
             key='name',
-            value=self.module.params.get('os'),
+            value=os_name,
             resource='os',
             use_cache=True
         )
 
+    def get_snapshot(self):
+        return self.query_resource_by_key(
+            key='description',
+            value=self.module.params.get('snapshot'),
+            resource='snapshot',
+            use_cache=True
+        )
+
     def get_ssh_keys(self):
         ssh_key_names = self.module.params.get('ssh_keys')
         if not ssh_key_names:
@@ -490,6 +507,11 @@ class AnsibleVultrServer(Vultr):
             'plan',
             'region',
         ]
+
+        snapshot_restore = self.module.params.get('snapshot') is not None
+        if snapshot_restore:
+            required_params.remove('os')
+
         self.module.fail_on_missing_params(required_params=required_params)
 
         self.result['changed'] = True
@@ -499,6 +521,7 @@ class AnsibleVultrServer(Vultr):
                 'VPSPLANID': self.get_plan().get('VPSPLANID'),
                 'FIREWALLGROUPID': self.get_firewall_group().get('FIREWALLGROUPID'),
                 'OSID': self.get_os().get('OSID'),
+                'SNAPSHOTID': self.get_snapshot().get('SNAPSHOTID'),
                 'label': self.module.params.get('name'),
                 'hostname': self.module.params.get('hostname'),
                 'SSHKEYID': ','.join([ssh_key['SSHKEYID'] for ssh_key in self.get_ssh_keys()]),
@@ -517,7 +540,7 @@ class AnsibleVultrServer(Vultr):
                 data=data
             )
             server = self._wait_for_state(key='status', state='active')
-            server = self._wait_for_state(state='running')
+            server = self._wait_for_state(state='running', timeout=3600 if snapshot_restore else 60)
         return server
 
     def _update_auto_backups_setting(self, server, start_server):
@@ -640,6 +663,10 @@ class AnsibleVultrServer(Vultr):
         return server, warned
 
     def _update_server(self, server=None, start_server=True):
+        # Wait for server to unlock if restoring
+        if server.get('os').strip() == 'Snapshot':
+            server = self._wait_for_state(key='server_status', state='ok', timeout=3600)
+
         # Update auto backups settings, stops server
         server = self._update_auto_backups_setting(server=server, start_server=start_server)
 
@@ -770,10 +797,10 @@ class AnsibleVultrServer(Vultr):
                 server = self._wait_for_state(state='running')
         return server
 
-    def _wait_for_state(self, key='power_status', state=None):
+    def _wait_for_state(self, key='power_status', state=None, timeout=60):
         time.sleep(1)
         server = self.get_server(refresh=True)
-        for s in range(0, 60):
+        for s in range(0, timeout):
             # Check for Truely if wanted state is None
             if state is None and server.get(key):
                 break
@@ -840,6 +867,7 @@ def main():
         name=dict(required=True, aliases=['label']),
         hostname=dict(),
         os=dict(),
+        snapshot=dict(),
         plan=dict(),
         force=dict(type='bool', default=False),
         notify_activate=dict(type='bool', default=False),