From 79f4dc91ae5c29bf6083ae7ad12d1001c3638170 Mon Sep 17 00:00:00 2001
From: Matous Mojzis <mmojzis@redhat.com>
Date: Tue, 28 Jan 2020 14:42:00 +0100
Subject: [PATCH] Added snapshot import functions

---
 wrapanapi/systems/ec2.py | 36 ++++++++++++++++++++++++++++++++++++
 1 file changed, 36 insertions(+)

diff --git a/wrapanapi/systems/ec2.py b/wrapanapi/systems/ec2.py
index b24b8a3f..4c859851 100644
--- a/wrapanapi/systems/ec2.py
+++ b/wrapanapi/systems/ec2.py
@@ -1471,3 +1471,39 @@ def set_sns_topic_target_for_all_cw_rules(self, topic_arn):
             return True
         except Exception:
             return False
+
+    def import_snapshot(self, s3bucket, s3key, format="vhd", description=None):
+        self.logger.info(
+            " Importing snapshot %s from %s bucket with description %s in %s started successfully.",
+            s3key, s3bucket, description, format
+        )
+        try:
+            result = self.ec2_connection.import_snapshot(
+                DiskContainer={
+                    'Description': description if description is not None else s3key,
+                    'Format': format,
+                    'UserBucket': {
+                        'S3Bucket': s3bucket,
+                        'S3Key': s3key
+                    }
+                }
+            )
+            task_id = result.get("ImportTaskId")
+            return task_id
+
+        except Exception:
+            self.logger.exception("Import of snapshot '%s' failed.", s3key)
+            return False
+
+    def get_import_snapshot_task(self, task_id):
+        result = self.ec2_connection.describe_import_snapshot_tasks(ImportTaskIds=[task_id])
+        result_task = result.get("ImportSnapshotTasks")
+        return result_task[0]
+
+    def get_snapshot_id_if_import_completed(self, task_id):
+        result = self.get_import_snapshot_task(task_id).get('SnapshotTaskDetail')
+        result_status = result.get("Status")
+        if result_status == 'completed':
+            return result.get("SnapshotId")
+        else:
+            return False