Zip/unzip files on archive/unarchive imports

1 files changed, 100 insertions, 3 deletions

diff --git a/ishtar_common/models_imports.py b/ishtar_common/models_imports.py
index 7a8a10bc2..5a3af1a05 100644
--- a/ishtar_common/models_imports.py
+++ b/ishtar_common/models_imports.py
@@ -19,6 +19,7 @@
 
 import csv
 import datetime
+import json
 import os
 import logging
 import shutil
@@ -26,9 +27,11 @@ import re
 import tempfile
 import zipfile
 
+from django.apps import apps
 from django.conf import settings
 from django.contrib.gis.db import models
 from django.core.exceptions import ValidationError
+from django.core.files import File
 from django.core.files.base import ContentFile
 from django.core.validators import validate_comma_separated_integer_list
 from django.db.models.base import ModelBase
@@ -965,6 +968,9 @@ class Import(models.Model):
     match_file = models.FileField(
         _("Match file"), upload_to="upload/imports/%Y/%m/", blank=True,
         null=True, max_length=255, help_text=max_size_help())
+    archive_file = models.FileField(
+        _("Archive file"), upload_to="upload/imports/%Y/%m/", blank=True,
+        null=True, max_length=255, help_text=max_size_help())
     state = models.CharField(_("State"), max_length=2, choices=IMPORT_STATE,
                              default='C')
     conservative_import = models.BooleanField(
@@ -1096,7 +1102,7 @@ class Import(models.Model):
         """
         Get available action relevant with the current status
         """
-        from ishtar_common.models import IshtarSiteProfile
+        IshtarSiteProfile = apps.get_model("ishtar_common", "IshtarSiteProfile")
         profile = IshtarSiteProfile.get_current_profile()
         actions = []
         if self.state == 'C':
@@ -1121,7 +1127,8 @@ class Import(models.Model):
                     actions.append(('CH', _("Check for changes")))
             actions.append(('AC', _("Archive")))
         if self.state == 'AC':
-            actions.append(('A', _("Unarchive")))
+            state = "FE" if self.error_file else "F"
+            actions.append((state, _("Unarchive")))
         actions.append(('D', _("Delete")))
         return actions
 
@@ -1357,10 +1364,94 @@ class Import(models.Model):
         if return_importer_and_data:
             return importer, data
 
+    def _unarchive(self):
+        if not self.archive_file:
+            return
+        with tempfile.TemporaryDirectory() as tmp_dir_name:
+            # extract the current archive
+            current_zip = zipfile.ZipFile(self.archive_file.path, 'r')
+            name_list = current_zip.namelist()
+            if "content.json" not in name_list:
+                return
+            for name in name_list:
+                current_zip.extract(name, tmp_dir_name)
+            current_zip.close()
+            content_name = os.path.join(tmp_dir_name, "content.json")
+            try:
+                with open(content_name, "r") as content:
+                    files = json.loads(content.read())
+            except (IOError, json.JSONDecodeError):
+                return
+            today = datetime.date.today()
+            for attr in files:
+                filename = files[attr]
+                full_filename = os.path.join(tmp_dir_name, filename)
+                with open(full_filename, "rb") as raw_file:
+                    getattr(self, attr).save(
+                        "upload/imports/{}/{:02d}/{}".format(
+                            today.year, today.month, filename),
+                        File(raw_file)
+                    )
+
+        os.remove(self.archive_file.path)
+        setattr(self, 'archive_file', None)
+        self.state = "FE" if self.error_file else "F"
+        self.save()
+        return True
+
+    def _archive(self):
+        file_attr = ["imported_file", "error_file", "result_file",
+                     "match_file"]
+        files = [
+            (k, getattr(self, k).path, getattr(self, k).name.split(os.sep)[-1])
+            for k in file_attr
+            if getattr(self, k)
+        ]
+        self._archive_pending = True
+        with tempfile.TemporaryDirectory() as tmpdir:
+            base_name = "{}.zip".format(slugify(self.name))
+            archive_name = os.path.join(tmpdir, base_name)
+            with zipfile.ZipFile(archive_name, "w") as current_zip:
+                zip_content = {}
+                for k, path, name in files:
+                    try:
+                        current_zip.write(path, arcname=name)
+                        zip_content[k] = name
+                    except OSError:
+                        pass
+                content_name = os.path.join(tmpdir, "content.json")
+                with open(content_name, "w") as content:
+                    content.write(json.dumps(zip_content))
+                current_zip.write(content_name, arcname="content.json")
+
+            today = datetime.date.today()
+            with open(archive_name, "rb", ) as raw_file:
+                self.archive_file.save(
+                    "upload/imports/{}/{:02d}/{}".format(
+                        today.year, today.month, base_name),
+                    File(raw_file)
+                )
+        IshtarSiteProfile = apps.get_model("ishtar_common", "IshtarSiteProfile")
+        profile = IshtarSiteProfile.get_current_profile()
+        if profile.delete_image_zip_on_archive:
+            file_attr.append("imported_images")
+        for attr in file_attr:
+            file_field = getattr(self, attr)
+            if file_field:
+                os.remove(file_field.path)
+                setattr(self, attr, None)
+        self.save()
+        self._archive_pending = False
+
     def archive(self):
         self.state = 'AC'
         self.end_date = datetime.datetime.now()
-        self.save()
+        self._archive()
+
+    def unarchive(self, state):
+        if not self._unarchive():
+            self.state = state
+            self.save()  # only save if no save previously
 
     def get_all_imported(self):
         imported = []
@@ -1370,6 +1461,12 @@ class Import(models.Model):
                          for obj in getattr(self, accessor).all()]
         return imported
 
+    def save(self, *args, **kwargs):
+        super(Import, self).save(*args, **kwargs)
+        if self.state == "AC" and not getattr(
+                self, "_archive_pending", False) and not self.archive_file:
+            self._archive()
+
 
 def pre_delete_import(sender, **kwargs):
     # deleted imported items when an import is delete