diff --git a/modoboa/admin/management/commands/subcommands/_export.py b/modoboa/admin/management/commands/subcommands/_export.py index 94beef93c..15ff8cf2a 100644 --- a/modoboa/admin/management/commands/subcommands/_export.py +++ b/modoboa/admin/management/commands/subcommands/_export.py @@ -6,14 +6,18 @@ from __future__ import unicode_literals import sys -from backports import csv - from django.core.management.base import BaseCommand +from django.utils import six from modoboa.core.extensions import exts_pool from modoboa.core.models import User from .... import models +if six.PY2: + from backports import csv +else: + import csv + class ExportCommand(BaseCommand): """Command class.""" diff --git a/modoboa/admin/management/commands/subcommands/_import.py b/modoboa/admin/management/commands/subcommands/_import.py index 0c607b604..a202fe22b 100644 --- a/modoboa/admin/management/commands/subcommands/_import.py +++ b/modoboa/admin/management/commands/subcommands/_import.py @@ -8,15 +8,23 @@ import io import os import progressbar -from backports import csv +from chardet.universaldetector import UniversalDetector from django.core.management.base import BaseCommand, CommandError +from django.db import transaction +from django.utils import six +from django.utils.translation import ugettext as _ from modoboa.core import models as core_models from modoboa.core.extensions import exts_pool from modoboa.lib.exceptions import Conflict from .... import signals +if six.PY2: + from backports import csv +else: + import csv + class ImportCommand(BaseCommand): """Command class.""" @@ -27,7 +35,7 @@ class ImportCommand(BaseCommand): def add_arguments(self, parser): """Add extra arguments to command.""" parser.add_argument( - "--sepchar", type=str, default=";", + "--sepchar", type=six.text_type, default=";", help="Separator used in file.") parser.add_argument( "--continue-if-exists", action="store_true", @@ -37,23 +45,25 @@ class ImportCommand(BaseCommand): "--crypt-password", action="store_true", default=False, help="Encrypt provided passwords.") parser.add_argument( - "files", type=str, nargs="+", help="CSV files to import.") + "files", type=six.text_type, nargs="+", help="CSV files to import.") - def _import(self, filename, options): + def _import(self, filename, options, encoding="utf-8"): """Import domains or identities.""" superadmin = ( core_models.User.objects.filter(is_superuser=True).first() ) - if not os.path.exists(filename): + if not os.path.isfile(filename): raise CommandError("File not found") - num_lines = sum(1 for line in open(filename) if line) + num_lines = sum( + 1 for line in io.open(filename, encoding=encoding) if line + ) pbar = progressbar.ProgressBar( widgets=[ progressbar.Percentage(), progressbar.Bar(), progressbar.ETA() ], maxval=num_lines ).start() - with io.open(filename, encoding="utf8") as f: + with io.open(filename, encoding=encoding, newline="") as f: reader = csv.reader(f, delimiter=options["sepchar"]) i = 0 for row in reader: @@ -82,4 +92,39 @@ class ImportCommand(BaseCommand): """Command entry point.""" exts_pool.load_all() for filename in options["files"]: - self._import(filename, options) + try: + with transaction.atomic(): + self._import(filename, options) + except CommandError as exc: + raise exc + except UnicodeDecodeError: + self.stdout.write(self.style.NOTICE( + _("CSV file is not encoded in UTF-8, attempting to guess " + "encoding") + )) + detector = UniversalDetector() + with io.open(filename, "rb") as fp: + for line in fp: + detector.feed(line) + if detector.done: + break + detector.close() + + self.stdout.write(self.style.NOTICE( + _("Reading CSV file using %(encoding)s encoding") % + detector.result + )) + try: + with transaction.atomic(): + self._import( + filename, options, + encoding=detector.result["encoding"] + ) + except UnicodeDecodeError as exc: + six.raise_from( + CommandError( + _("Unable to decode CSV file using %(encoding)s " + "encoding") % detector.result + ), + exc + ) diff --git a/modoboa/admin/tests/test_data/import_domains.csv b/modoboa/admin/tests/test_data/import_domains.csv new file mode 100644 index 000000000..d879f8a16 --- /dev/null +++ b/modoboa/admin/tests/test_data/import_domains.csv @@ -0,0 +1,2 @@ +domain;domain1.com;1000;100;True +domain;dómain2.com;2000;200;True diff --git a/modoboa/admin/tests/test_data/import_domains_duplicates.csv b/modoboa/admin/tests/test_data/import_domains_duplicates.csv new file mode 100644 index 000000000..c5395f364 --- /dev/null +++ b/modoboa/admin/tests/test_data/import_domains_duplicates.csv @@ -0,0 +1,3 @@ +domain;domain1.com;1000;100;True +domain;dómain2.com;2000;200;True +domain;dómain2.com;2000;200;True diff --git a/modoboa/admin/tests/test_import_.py b/modoboa/admin/tests/test_import_.py index 148b288a3..c5d49760c 100644 --- a/modoboa/admin/tests/test_import_.py +++ b/modoboa/admin/tests/test_import_.py @@ -2,9 +2,13 @@ from __future__ import unicode_literals +import os + import dns.resolver from django.core.files.base import ContentFile +from django.core.management import call_command +from django.core.management.base import CommandError from django.urls import reverse from modoboa.core import factories as core_factories @@ -353,3 +357,63 @@ dómain; dómªin2.com; 1000; 100; True self.assertEqual(dom.default_mailbox_quota, 100) self.assertTrue(dom.enabled) self.assertTrue(admin.is_owner(dom)) + + def test_import_command_missing_file(self): + with mock.patch("os.path.isfile") as mock_isfile: + mock_isfile.return_value = False + with self.assertRaises(CommandError) as cm: + call_command("modo", "import", "test.csv") + ex_message = cm.exception.messages + self.assertEqual(ex_message, "File not found") + + def test_import_command(self): + test_file = os.path.join( + os.path.dirname(__file__), + "test_data/import_domains.csv" + ) + call_command("modo", "import", test_file) + + admin = User.objects.get(username="admin") + + dom = Domain.objects.get(name="domain1.com") + self.assertEqual(dom.quota, 1000) + self.assertEqual(dom.default_mailbox_quota, 100) + self.assertTrue(dom.enabled) + self.assertTrue(admin.is_owner(dom)) + + dom = Domain.objects.get(name="dómain2.com") + self.assertEqual(dom.quota, 2000) + self.assertEqual(dom.default_mailbox_quota, 200) + self.assertTrue(dom.enabled) + self.assertTrue(admin.is_owner(dom)) + + def test_import_command_duplicates(self): + test_file = os.path.join( + os.path.dirname(__file__), + "test_data/import_domains_duplicates.csv" + ) + with self.assertRaises(CommandError) as cm: + call_command("modo", "import", test_file) + ex_message = cm.exception.messages + self.assertTrue(ex_message.startswith("Object already exists: ")) + + def test_import_command_duplicates_continue(self): + test_file = os.path.join( + os.path.dirname(__file__), + "test_data/import_domains_duplicates.csv" + ) + call_command("modo", "import", "--continue-if-exists", test_file) + + admin = User.objects.get(username="admin") + + dom = Domain.objects.get(name="domain1.com") + self.assertEqual(dom.quota, 1000) + self.assertEqual(dom.default_mailbox_quota, 100) + self.assertTrue(dom.enabled) + self.assertTrue(admin.is_owner(dom)) + + dom = Domain.objects.get(name="dómain2.com") + self.assertEqual(dom.quota, 2000) + self.assertEqual(dom.default_mailbox_quota, 200) + self.assertTrue(dom.enabled) + self.assertTrue(admin.is_owner(dom)) diff --git a/modoboa/admin/views/export.py b/modoboa/admin/views/export.py index 9773b3c6a..281d864d3 100644 --- a/modoboa/admin/views/export.py +++ b/modoboa/admin/views/export.py @@ -4,7 +4,6 @@ from __future__ import unicode_literals -from backports import csv from rfc6266 import build_header from django.contrib.auth.decorators import ( @@ -19,6 +18,11 @@ from django.utils.translation import ugettext as _ from ..forms import ExportDomainsForm, ExportIdentitiesForm from ..lib import get_domains, get_identities +if six.PY2: + from backports import csv +else: + import csv + def _export(content, filename): """Export a csv file's content diff --git a/modoboa/admin/views/import_.py b/modoboa/admin/views/import_.py index 4544d47e2..025d9d0e3 100644 --- a/modoboa/admin/views/import_.py +++ b/modoboa/admin/views/import_.py @@ -6,7 +6,6 @@ from __future__ import unicode_literals import io -from backports import csv from reversion import revisions as reversion from django.contrib.auth.decorators import ( @@ -15,6 +14,7 @@ from django.contrib.auth.decorators import ( from django.db import transaction from django.shortcuts import render from django.urls import reverse +from django.utils import six from django.utils.encoding import smart_text from django.utils.translation import ugettext as _ @@ -22,6 +22,11 @@ from modoboa.lib.exceptions import Conflict, ModoboaException from .. import signals from ..forms import ImportDataForm, ImportIdentitiesForm +if six.PY2: + from backports import csv +else: + import csv + @reversion.create_revision() def importdata(request, formclass=ImportDataForm): diff --git a/requirements.txt b/requirements.txt index 632848c14..20965afed 100644 --- a/requirements.txt +++ b/requirements.txt @@ -26,5 +26,5 @@ pytz requests rfc6266 lxml -backports.csv +backports.csv; python_version < '3' chardet