From aefa886ce3bef72eccfc9e808bfaf645961104d0 Mon Sep 17 00:00:00 2001
From: Nicolas ROCHE <nroche@entrouvert.com>
Date: Sat, 1 Feb 2020 02:21:09 +0100
Subject: [PATCH] assets: add import_assets management command (#39425)

---
 combo/apps/assets/management/__init__.py      |  0
 .../assets/management/commands/__init__.py    |  0
 .../management/commands/export_assets.py      | 40 ++++++++
 .../management/commands/import_assets.py      | 47 +++++++++
 combo/apps/assets/utils.py                    | 38 ++++++++
 combo/apps/assets/views.py                    | 18 +---
 tests/test_assets.py                          | 96 ++++++++++++++++++-
 7 files changed, 223 insertions(+), 16 deletions(-)
 create mode 100644 combo/apps/assets/management/__init__.py
 create mode 100644 combo/apps/assets/management/commands/__init__.py
 create mode 100644 combo/apps/assets/management/commands/export_assets.py
 create mode 100644 combo/apps/assets/management/commands/import_assets.py
 create mode 100644 combo/apps/assets/utils.py

diff --git a/combo/apps/assets/management/__init__.py b/combo/apps/assets/management/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/combo/apps/assets/management/commands/__init__.py b/combo/apps/assets/management/commands/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/combo/apps/assets/management/commands/export_assets.py b/combo/apps/assets/management/commands/export_assets.py
new file mode 100644
index 00000000..3efc5c3d
--- /dev/null
+++ b/combo/apps/assets/management/commands/export_assets.py
@@ -0,0 +1,40 @@
+# combo - content management system
+# Copyright (C) 2020  Entr'ouvert
+#
+# This program is free software: you can redistribute it and/or modify it
+# under the terms of the GNU Affero General Public License as published
+# by the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# 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 sys
+import tarfile
+
+from django.core.management.base import BaseCommand, CommandError
+
+from combo.apps.assets.utils import export_assets
+
+class Command(BaseCommand):
+    help = 'Export assets'
+
+    def add_arguments(self, parser):
+        parser.add_argument(
+                '--output', metavar='FILENAME', default='assets.tar',
+                help='name of a file to write output to')
+
+    def handle(self, *args, **options):
+        if options['output'] != '-':
+            try:
+                assets_file = tarfile.open(options['output'], 'w')
+            except IOError as e:
+                raise CommandError(e)
+        else:
+            assets_file = tarfile.open(mode='w|', fileobj=sys.stdout.buffer)
+        export_assets(assets_file)
diff --git a/combo/apps/assets/management/commands/import_assets.py b/combo/apps/assets/management/commands/import_assets.py
new file mode 100644
index 00000000..105fc9f3
--- /dev/null
+++ b/combo/apps/assets/management/commands/import_assets.py
@@ -0,0 +1,47 @@
+# combo - content management system
+# Copyright (C) 2020  Entr'ouvert
+#
+# This program is free software: you can redistribute it and/or modify it
+# under the terms of the GNU Affero General Public License as published
+# by the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# 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/>.
+
+from io import BytesIO
+import sys
+import tarfile
+
+from django.core.management.base import BaseCommand, CommandError
+
+from combo.apps.assets.views import import_assets
+
+class Command(BaseCommand):
+    help = 'Import an exported asset file'
+
+    def add_arguments(self, parser):
+        parser.add_argument('filename', metavar='FILENAME', type=str,
+                help='name of file to import')
+        parser.add_argument(
+                '--overwrite', action='store_true', default=False,
+                help='re-import assets already there')
+
+    def handle(self, filename, *args, **options):
+        if filename == '-':
+            fd = BytesIO(sys.stdin.buffer.read())
+        else:
+            try:
+                fd = open(filename, 'rb')
+            except IOError as e:
+                raise CommandError(e)
+        try:
+            assets = tarfile.open(fileobj=fd)
+        except tarfile.TarError as e:
+            raise CommandError(e)
+        import_assets(assets, options['overwrite'])
diff --git a/combo/apps/assets/utils.py b/combo/apps/assets/utils.py
new file mode 100644
index 00000000..9e4227dc
--- /dev/null
+++ b/combo/apps/assets/utils.py
@@ -0,0 +1,38 @@
+# combo - content management system
+# Copyright (C) 2020  Entr'ouvert
+#
+# This program is free software: you can redistribute it and/or modify it
+# under the terms of the GNU Affero General Public License as published
+# by the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# 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 os
+
+from django.core.files.storage import default_storage
+
+
+def import_assets(assets_tarfile, overwrite=False):
+    media_prefix = default_storage.path('')
+    for tarinfo in assets_tarfile.getmembers():
+        filepath = default_storage.path(tarinfo.name)
+        if not overwrite and os.path.exists(filepath):
+            continue
+        assets_tarfile.extract(tarinfo, path=media_prefix)
+
+
+def export_assets(assets_tarfile):
+    media_prefix = default_storage.path('')
+    for basedir, dirnames, filenames in os.walk(media_prefix):
+        for filename in filenames:
+            assets_tarfile.add(
+                    os.path.join(basedir, filename),
+                    os.path.join(basedir, filename)[len(media_prefix):])
+    assets_tarfile.close()
diff --git a/combo/apps/assets/views.py b/combo/apps/assets/views.py
index 991a9e00..07810868 100644
--- a/combo/apps/assets/views.py
+++ b/combo/apps/assets/views.py
@@ -32,6 +32,7 @@ from django.views.generic import TemplateView, ListView, FormView
 import ckeditor
 from sorl.thumbnail.shortcuts import get_thumbnail
 
+from combo.apps.assets.utils import import_assets, export_assets
 from combo.data.models import CellBase
 
 from .forms import AssetUploadForm, AssetsImportForm
@@ -267,16 +268,11 @@ class AssetsImport(FormView):
     def form_valid(self, form):
         overwrite = form.cleaned_data.get('overwrite')
         try:
-            assets = tarfile.open(fileobj=form.cleaned_data['assets_file'])
+            assets_file = tarfile.open(fileobj=form.cleaned_data['assets_file'])
         except tarfile.TarError:
             messages.error(self.request, _('The assets file is not valid.'))
             return super(AssetsImport, self).form_valid(form)
-        media_prefix = default_storage.path('')
-        for tarinfo in assets.getmembers():
-            filepath = default_storage.path(tarinfo.name)
-            if not overwrite and os.path.exists(filepath):
-                continue
-            assets.extract(tarinfo, path=media_prefix)
+        import_assets(assets_file, overwrite)
         messages.success(self.request, _('The assets file has been imported.'))
         return super(AssetsImport, self).form_valid(form)
 
@@ -286,13 +282,7 @@ assets_import = AssetsImport.as_view()
 def assets_export(request, *args, **kwargs):
     fd = BytesIO()
     assets_file = tarfile.open('assets.tar', 'w', fileobj=fd)
-    media_prefix = default_storage.path('')
-    for basedir, dirnames, filenames in os.walk(media_prefix):
-        for filename in filenames:
-            assets_file.add(
-                    os.path.join(basedir, filename),
-                    os.path.join(basedir, filename)[len(media_prefix):])
-    assets_file.close()
+    export_assets(assets_file)
     return HttpResponse(fd.getvalue(), content_type='application/x-tar')
 
 
diff --git a/tests/test_assets.py b/tests/test_assets.py
index 2cc051d9..b46527aa 100644
--- a/tests/test_assets.py
+++ b/tests/test_assets.py
@@ -1,15 +1,43 @@
 # -*- coding: utf-8 -*-
 
 import base64
+import io
+import os
+import pytest
+import sys
+import tarfile
 
-from django.core.urlresolvers import reverse
+from combo.apps.assets.utils import export_assets, import_assets
+from django.core.files.storage import default_storage
 
-import pytest
+from django.core.files import File
+from django.core.management import call_command, CommandError
+from django.core.urlresolvers import reverse
+from django.utils.six import BytesIO
 
 from combo.apps.assets.models import Asset
 
 pytestmark = pytest.mark.django_db
 
+
+@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 count_asset_files():
+    nb_assets = 0
+    media_prefix = default_storage.path('')
+    for basedir, dirnames, filenames in os.walk(media_prefix):
+        nb_assets += len(filenames)
+    return nb_assets
+
+def remove_asset_files():
+    media_prefix = default_storage.path('')
+    for basedir, dirnames, filenames in os.walk(media_prefix):
+        for filename in filenames:
+            os.remove('%s/%s' % (basedir, filename))
+
 def test_asset_set_api(app, john_doe):
     app.authorization = ('Basic', (john_doe.username, john_doe.username))
     resp = app.post_json(reverse('api-assets-set', kwargs={'key': 'plop'}), params={
@@ -42,3 +70,67 @@ def test_asset_set_api(app, john_doe):
             }
             }, status=400)
         assert resp.json.get('err') == 1
+
+def test_import_export_assets(app, some_assets):
+    assert Asset.objects.count() == 2
+    assert count_asset_files() == 2
+    fd = BytesIO()
+    assets_file = tarfile.open(mode='w', fileobj=fd)
+    export_assets(assets_file)
+    tar_bytes = fd.getvalue()
+
+    path = default_storage.path('')
+    os.remove('%s/assets/test.png' % path)
+    open('%s/assets/test2.png' % path, 'w').write('foo')
+    fd = BytesIO(tar_bytes)
+    assert count_asset_files() == 1
+    assets_file = tarfile.open(fileobj=fd)
+    import_assets(assets_file)
+    assert count_asset_files() == 2
+    assert open('%s/assets/test.png' % path, 'r').read() == 'test'
+    assert open('%s/assets/test2.png' % path, 'r').read() == 'foo'
+
+    fd = BytesIO(tar_bytes)
+    assets_file = tarfile.open(fileobj=fd)
+    import_assets(assets_file, True)
+    assert open('%s/assets/test2.png' % path, 'r').read() == 'test2'
+    remove_asset_files()
+
+def test_import_export_assets_cmd(some_assets, tmpdir):
+    assert Asset.objects.count() == 2
+    assert count_asset_files() == 2
+    call_command('export_assets')
+    path = default_storage.path('')
+    os.remove('%s/assets/test.png' % path)
+    open('%s/assets/test2.png' % path, 'w').write('foo')
+    assert count_asset_files() == 1
+    call_command('import_assets', 'assets.tar')
+    assert count_asset_files() == 2
+    assert open('%s/assets/test.png' % path, 'r').read() == 'test'
+    assert open('%s/assets/test2.png' % path, 'r').read() == 'foo'
+    remove_asset_files()
+
+    with pytest.raises(CommandError, match=r'No such file or directory'):
+        call_command('export_assets', '--output', '%s/noway/foo.tar' % tmpdir)
+    with pytest.raises(CommandError, match=r'No such file or directory'):
+        call_command('import_assets', '%s/foo.tar' % tmpdir)
+    open('%s/foo.tar' % tmpdir, 'w').write('foo')
+    with pytest.raises(CommandError, match=r'could not be opened successfully'):
+        call_command('import_assets', '%s/foo.tar' % tmpdir)
+
+def test_import_export_assets_cmd_flow(some_assets, capsysbinary, monkeypatch):
+    assert Asset.objects.count() == 2
+    assert count_asset_files() == 2
+    call_command('export_assets', output='-')
+    captured = capsysbinary.readouterr()
+    path = default_storage.path('')
+    os.remove('%s/assets/test.png' % path)
+    open('%s/assets/test2.png' % path, 'w').write('foo')
+    assert count_asset_files() == 1
+    with capsysbinary.disabled():
+        monkeypatch.setattr('sys.stdin', io.TextIOWrapper(BytesIO(captured.out)))
+        call_command('import_assets', '-', '--overwrite')
+    assert count_asset_files() == 2
+    assert open('%s/assets/test.png' % path, 'r').read() == 'test'
+    assert open('%s/assets/test2.png' % path, 'r').read() == 'test2'
+    remove_asset_files()
-- 
2.24.1

