From 4812f19546e2f0e28fe2513dc933cf24b562d46b Mon Sep 17 00:00:00 2001
From: Nicolas ROCHE <nroche@entrouvert.com>
Date: Mon, 13 Jul 2020 17:45:49 +0200
Subject: [PATCH 2/3] commands: add tar format for site export/import (#39426)

---
 combo/data/management/commands/export_site.py | 26 ++++++--
 combo/data/management/commands/import_site.py | 35 ++++++++--
 combo/data/utils.py                           | 39 ++++++++++-
 tests/test_import_export.py                   | 66 ++++++++++++++++++-
 tests/test_pages.py                           |  4 +-
 5 files changed, 153 insertions(+), 17 deletions(-)

diff --git a/combo/data/management/commands/export_site.py b/combo/data/management/commands/export_site.py
index cbe95f36..d9f02f2f 100644
--- a/combo/data/management/commands/export_site.py
+++ b/combo/data/management/commands/export_site.py
@@ -12,26 +12,40 @@
 # GNU Affero General Public License for more details.
 #
 # You should have received a copy of the GNU Affero General Public License
 # along with this program.  If not, see <http://www.gnu.org/licenses/>.
 
 import json
 import sys
 
-from django.core.management.base import BaseCommand
+from django.core.management.base import BaseCommand, CommandError
+from django.utils.translation import ugettext_lazy as _
 
-from combo.data.utils import export_site
+from combo.data.utils import export_site, export_site_tar
 
 class Command(BaseCommand):
     help = 'Export the site'
 
     def add_arguments(self, parser):
         parser.add_argument(
                 '--output', metavar='FILE', default=None,
                 help='name of a file to write output to')
+        parser.add_argument(
+                '--format-json', action='store_true', default=False,
+                help='use JSON format with no asset files')
 
     def handle(self, *args, **options):
-        if options['output'] and options['output'] != '-':
-            output = open(options['output'], 'w')
+        if options['format_json']:
+            if options['output'] and options['output'] != '-':
+                output = open(options['output'], 'w')
+            else:
+                output = sys.stdout
+            json.dump(export_site(), output, indent=2)
         else:
-            output = sys.stdout
-        json.dump(export_site(), output, indent=2)
+            if options['output'] and options['output'] != '-':
+                try:
+                    output = open(options['output'], 'wb')
+                except IOError as e:
+                    raise CommandError(e)
+                export_site_tar(output)
+            else:
+                raise CommandError(_('TAR format require output filename parameter'))
diff --git a/combo/data/management/commands/import_site.py b/combo/data/management/commands/import_site.py
index 509a5377..6f2d0bd9 100644
--- a/combo/data/management/commands/import_site.py
+++ b/combo/data/management/commands/import_site.py
@@ -11,38 +11,59 @@
 # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 # GNU Affero General Public License for more details.
 #
 # You should have received a copy of the GNU Affero General Public License
 # along with this program.  If not, see <http://www.gnu.org/licenses/>.
 
 import json
 import sys
+import tarfile
 
 from django.core.management.base import BaseCommand, CommandError
-from django.utils.encoding import force_text
 
-from combo.data.utils import import_site, MissingGroups
+from combo.data.utils import import_site, import_site_tar, ImportSiteError
 
 class Command(BaseCommand):
     help = 'Import an exported site'
 
     def add_arguments(self, parser):
         parser.add_argument('filename', metavar='FILENAME', type=str,
                 help='name of file to import')
         parser.add_argument(
                 '--clean', action='store_true', default=False,
                 help='Clean site before importing')
         parser.add_argument(
                 '--if-empty', action='store_true', default=False,
                 help='Import only if site is empty')
+        parser.add_argument(
+                '--overwrite', action='store_true', default=False,
+                help='Overwrite asset files')
 
     def handle(self, filename, *args, **options):
         if filename == '-':
+            format = 'json'
             fd = sys.stdin
         else:
-            fd = open(filename)
+            try:
+                fd = open(filename, 'rb')
+            except IOError as e:
+                raise CommandError(e)
+            try:
+                tarfile.open(mode='r', fileobj=fd)
+            except tarfile.TarError as e:
+                format = 'json'
+                fd = open(filename, 'r')
+            else:
+                format = 'tar'
+                fd = open(filename, 'rb')
         try:
-            import_site(json.load(fd),
-                        if_empty=options['if_empty'],
-                        clean=options['clean'])
-        except MissingGroups as e:
+            if format == 'json':
+                import_site(json.load(fd),
+                            if_empty=options['if_empty'],
+                            clean=options['clean'])
+            else:
+                import_site_tar(fd,
+                                if_empty=options['if_empty'],
+                                clean=options['clean'],
+                                overwrite=options['overwrite'])
+        except ImportSiteError as e:
             raise CommandError(e)
diff --git a/combo/data/utils.py b/combo/data/utils.py
index 93365ffa..9198bd1e 100644
--- a/combo/data/utils.py
+++ b/combo/data/utils.py
@@ -9,30 +9,38 @@
 # This program is distributed in the hope that it will be useful,
 # but WITHOUT ANY WARRANTY; without even the implied warranty of
 # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 # GNU Affero General Public License for more details.
 #
 # You should have received a copy of the GNU Affero General Public License
 # along with this program.  If not, see <http://www.gnu.org/licenses/>.
 
+import json
+import tarfile
+
 from django.contrib.auth.models import Group
 from django.db import transaction
 from django.utils import six
 from django.utils.encoding import python_2_unicode_compatible
 from django.utils.translation import ugettext_lazy as _
 
 from combo.apps.assets.models import Asset
+from combo.apps.assets.utils import (add_tar_content, clean_assets_files,
+        untar_assets_files, tar_assets_files)
 from combo.apps.maps.models import MapLayer
 from combo.apps.pwa.models import PwaSettings, PwaNavigationEntry
 from .models import Page
 
 
+class ImportSiteError(Exception):
+    pass
+
 @python_2_unicode_compatible
-class MissingGroups(Exception):
+class MissingGroups(ImportSiteError):
     def __init__(self, names):
         self.names = names
 
     def __str__(self):
         return _('Missing groups: %s') % ', '.join(self.names)
 
 
 def export_site():
@@ -80,8 +88,37 @@ def import_site(data, if_empty=False, clean=False, request=None):
 
         MapLayer.load_serialized_objects(data.get('map-layers') or [])
         Asset.load_serialized_objects(data.get('assets') or [])
         Page.load_serialized_pages(data.get('pages') or [], request=request)
 
         if data.get('pwa'):
             PwaSettings.load_serialized_settings(data['pwa'].get('settings'))
             PwaNavigationEntry.load_serialized_objects(data['pwa'].get('navigation'))
+
+
+def export_site_tar(fd):
+    tar = tarfile.open(mode='w', fileobj=fd)
+    data = export_site()
+    del data['assets']
+    add_tar_content(tar, '_site.json', json.dumps(data, indent=2))
+    tar_assets_files(tar)
+    tar.close()
+
+
+def import_site_tar(fd, if_empty=False, clean=False, overwrite=False, request=None):
+    tar = tarfile.open(mode='r', fileobj=fd)
+    try:
+        tarinfo = tar.getmember('_site.json')
+    except KeyError:
+        raise ImportSiteError(_('TAR file should provide _site.json file'))
+
+    if if_empty and (Page.objects.count() or MapLayer.objects.count()):
+        return
+
+    if clean:
+        clean_assets_files()
+
+    json_site = tar.extractfile(tarinfo).read()
+    data = json.loads(json_site.decode('utf-8'))
+    data.update(untar_assets_files(tar, overwrite=overwrite))
+    import_site(data, if_empty=if_empty, clean=clean, request=request)
+    tar.close()
diff --git a/tests/test_import_export.py b/tests/test_import_export.py
index e8ab90c4..23117738 100644
--- a/tests/test_import_export.py
+++ b/tests/test_import_export.py
@@ -1,25 +1,28 @@
 import base64
 import datetime
 import json
 import os
 import shutil
 import sys
+import tarfile
 import tempfile
 
 import pytest
 from django.contrib.auth.models import Group
 from django.core.files import File
+from django.core.files.storage import default_storage
 from django.core.management import call_command
 from django.core.management.base import CommandError
 from django.utils.encoding import force_bytes, force_text
 from django.utils.six import BytesIO, StringIO
 
 from combo.apps.assets.models import Asset
+from combo.apps.assets.utils import clean_assets_files
 from combo.apps.gallery.models import Image, GalleryCell
 from combo.apps.maps.models import MapLayer, Map, MapLayerOptions
 from combo.apps.pwa.models import PwaSettings, PwaNavigationEntry
 from combo.data.models import Page, TextCell
 from combo.data.utils import export_site, import_site, MissingGroups
 
 pytestmark = pytest.mark.django_db
 
@@ -41,17 +44,17 @@ def some_map_layers():
 @pytest.fixture
 def some_assets():
     Asset(key='banner', asset=File(BytesIO(b'test'), 'test.png')).save()
     Asset(key='favicon', asset=File(BytesIO(b'test2'), 'test2.png')).save()
 
 def get_output_of_command(command, *args, **kwargs):
     old_stdout = sys.stdout
     output = sys.stdout = StringIO()
-    call_command(command, *args, **kwargs)
+    call_command(command, format_json=True, *args, **kwargs)
     sys.stdout = old_stdout
     return output.getvalue()
 
 
 def test_import_export(app, some_data):
     output = get_output_of_command('export_site')
     assert len(json.loads(output)['pages']) == 3
     import_site(data={}, clean=True)
@@ -329,8 +332,69 @@ def test_import_export_extra_fields(app, some_data):
         if page['fields']['slug'] == 'one':
             page['fields']['extra_field_not_in_model'] = True
         elif page['fields']['slug'] == 'three':
             page['cells'][0]['fields']['extra_field_not_in_model'] = True
 
     import_site(site_export)
     assert Page.objects.count() == 3
     assert TextCell.objects.count() == 1
+
+
+def test_import_export_tar(tmpdir, some_assets):
+    filename = os.path.join(str(tmpdir), 'file.tar')
+
+    # build import having some_assets fixtures assets: banner and favicon
+    call_command('export_site', '--output', filename)
+
+    def populate_site():
+        Page.objects.all().delete()
+        Asset.objects.all().delete()
+        clean_assets_files()
+        Page.objects.create(title='One', slug='one')
+        Asset(key='banner', asset=File(BytesIO(b'original content'), 'test.png')).save()
+        Asset(key='logo', asset=File(BytesIO(b'logo'), 'logo.png')).save()
+
+    populate_site()
+    call_command('import_site', filename)  # default behaviour
+    assert Page.objects.count() == 1
+    assert Asset.objects.count() == 3
+    Asset.objects.get(key='banner').asset.name == 'assets/test.png'
+    assert open('%s/assets/test.png' % default_storage.path('')).read() == 'original content'
+
+    populate_site()
+    call_command('import_site', filename, '--overwrite')
+    assert Page.objects.count() == 1
+    assert Asset.objects.count() == 3
+    Asset.objects.get(key='banner').asset.name == 'assets/test.png'
+    assert open('%s/assets/test.png' % default_storage.path('')).read() == 'test'
+
+    populate_site()
+    call_command('import_site', filename, '--if-empty')
+    assert Page.objects.count() == 1
+    assert Asset.objects.count() == 2
+    Asset.objects.get(key='banner').asset.name == 'assets/test3.png'
+    assert open('%s/assets/test.png' % default_storage.path('')).read() == 'original content'
+    Asset.objects.get(key='logo').asset.name == 'assets/logo.png'
+    assert os.path.isfile('%s/assets/logo.png' % default_storage.path(''))
+
+    populate_site()
+    call_command('import_site', filename, '--clean')
+    assert Page.objects.count() == 0
+    assert Asset.objects.count() == 2
+    Asset.objects.get(key='banner').asset.name == 'assets/test.png'
+    assert open('%s/assets/test.png' % default_storage.path('')).read() == 'test'
+    assert not Asset.objects.filter(key='logo')
+    assert not os.path.isfile('%s/assets/logo.png' % default_storage.path(''))
+
+    # error cases
+    with pytest.raises(CommandError, match=r'No such file or directory'):
+        call_command('export_site', '--output', '%s/noway/foo.tar' % tmpdir)
+
+    with pytest.raises(CommandError, match='TAR format require output filename parameter'):
+        call_command('export_site', '--output', '-')
+
+    with pytest.raises(CommandError, match=r'No such file or directory'):
+        call_command('import_site', '%s/noway/foo.tar' % tmpdir)
+
+    tarfile.open(filename, 'w').close()  # empty tar file
+    with pytest.raises(CommandError, match=r'TAR file should provide _site.json file'):
+        call_command('import_site', filename)
diff --git a/tests/test_pages.py b/tests/test_pages.py
index 6f623ed4..b6eeaef1 100644
--- a/tests/test_pages.py
+++ b/tests/test_pages.py
@@ -336,23 +336,23 @@ def test_import_export_management_commands():
     cell = TextCell(page=page2, text='bar', order=0, placeholder='content')
     cell.save()
 
     export_filename = os.path.join(settings.MEDIA_ROOT, 'site-export.json')
     if os.path.exists(export_filename):
         os.unlink(export_filename)
 
     cmd = ExportSiteCommand()
-    cmd.handle(output=export_filename)
+    cmd.handle(output=export_filename, format_json=True)
     assert os.path.exists(export_filename)
 
     stdout = sys.stdout
     try:
         sys.stdout = StringIO()
-        cmd.handle(output='-')
+        cmd.handle(output='-', format_json=True)
         assert sys.stdout.getvalue() == open(export_filename).read()
     finally:
         sys.stdout = stdout
 
     Page.objects.all().delete()
 
     cmd = ImportSiteCommand()
     cmd.handle(filename=export_filename, if_empty=False, clean=False)
-- 
2.27.0

