From 410bb8ad8979dff8b19a7ef5e7eb84a5d4e1936e Mon Sep 17 00:00:00 2001
From: Tushar912 <tushar.912u@gmail.com>
Date: Sat, 19 Dec 2020 23:14:53 +0530
Subject: [PATCH] imports: Add better checking for subdomains.

Add a `--allow-reserved-subdomain` flag which allows creation of
reserved keyword domains.  This also always enforces that the domain
is not in use, which was removed in 0258d7d.

Fixes #16924.
---
 zerver/forms.py                      |  9 ++++-----
 zerver/management/commands/import.py | 15 ++++++++++++++-
 zerver/tests/test_signup.py          | 28 ++++++++++++++++++++++------
 3 files changed, 40 insertions(+), 12 deletions(-)

diff --git a/zerver/forms.py b/zerver/forms.py
index 737ac995ad..93a195263e 100644
--- a/zerver/forms.py
+++ b/zerver/forms.py
@@ -65,7 +65,7 @@ def email_is_not_mit_mailing_list(email: str) -> None:
             else:
                 raise AssertionError("Unexpected DNS error")
 
-def check_subdomain_available(subdomain: str, from_management_command: bool=False) -> None:
+def check_subdomain_available(subdomain: str, allow_reserved_subdomain: bool=False) -> None:
     error_strings = {
         'too short': _("Subdomain needs to have length 3 or greater."),
         'extremal dash': _("Subdomain cannot start or end with a '-'."),
@@ -80,12 +80,11 @@ def check_subdomain_available(subdomain: str, from_management_command: bool=Fals
         raise ValidationError(error_strings['extremal dash'])
     if not re.match('^[a-z0-9-]*$', subdomain):
         raise ValidationError(error_strings['bad character'])
-    if from_management_command:
-        return
     if len(subdomain) < 3:
         raise ValidationError(error_strings['too short'])
-    if is_reserved_subdomain(subdomain) or \
-       Realm.objects.filter(string_id=subdomain).exists():
+    if Realm.objects.filter(string_id=subdomain).exists():
+        raise ValidationError(error_strings['unavailable'])
+    if is_reserved_subdomain(subdomain) and not allow_reserved_subdomain:
         raise ValidationError(error_strings['unavailable'])
 
 class RegistrationForm(forms.Form):
diff --git a/zerver/management/commands/import.py b/zerver/management/commands/import.py
index 73d68b2f99..0029e2e4d9 100644
--- a/zerver/management/commands/import.py
+++ b/zerver/management/commands/import.py
@@ -4,6 +4,7 @@ import subprocess
 from typing import Any
 
 from django.conf import settings
+from django.core.exceptions import ValidationError
 from django.core.management import call_command
 from django.core.management.base import BaseCommand, CommandError, CommandParser
 
@@ -26,6 +27,10 @@ import a database dump from one or more JSON files."""
                             action="store_true",
                             help='Import into an existing nonempty database.')
 
+        parser.add_argument('--allow-reserved-subdomain',
+                            action="store_true",
+                            help='Allow use of reserved subdomains')
+
         parser.add_argument('subdomain', metavar='<subdomain>',
                             help="Subdomain")
 
@@ -55,7 +60,15 @@ import a database dump from one or more JSON files."""
         elif options["import_into_nonempty"]:
             print("NOTE: The argument 'import_into_nonempty' is now the default behavior.")
 
-        check_subdomain_available(subdomain, from_management_command=True)
+        allow_reserved_subdomain = False
+
+        if options["allow_reserved_subdomain"]:
+            allow_reserved_subdomain = True
+
+        try:
+            check_subdomain_available(subdomain, allow_reserved_subdomain)
+        except ValidationError:
+            raise CommandError("Subdomain reserved: pass --allow-reserved-subdomain to use.")
 
         paths = []
         for path in options['export_paths']:
diff --git a/zerver/tests/test_signup.py b/zerver/tests/test_signup.py
index 1ff426d5aa..30dd3f1071 100644
--- a/zerver/tests/test_signup.py
+++ b/zerver/tests/test_signup.py
@@ -2634,13 +2634,29 @@ class RealmCreationTest(ZulipTestCase):
         self.assert_not_in_success_response(["unavailable"], result)
 
     def test_subdomain_check_management_command(self) -> None:
-        # Short names should work
-        check_subdomain_available('aa', from_management_command=True)
-        # So should reserved ones
-        check_subdomain_available('zulip', from_management_command=True)
-        # malformed names should still not
+        # Short names should not work, even with the flag
         with self.assertRaises(ValidationError):
-            check_subdomain_available('-ba_d-', from_management_command=True)
+            check_subdomain_available('aa')
+        with self.assertRaises(ValidationError):
+            check_subdomain_available('aa', allow_reserved_subdomain=True)
+
+        # Malformed names should never work
+        with self.assertRaises(ValidationError):
+            check_subdomain_available('-ba_d-')
+        with self.assertRaises(ValidationError):
+            check_subdomain_available('-ba_d-', allow_reserved_subdomain=True)
+
+        with patch('zerver.lib.name_restrictions.is_reserved_subdomain', return_value = False):
+            # Existing realms should never work even if they are not reserved keywords
+            with self.assertRaises(ValidationError):
+                check_subdomain_available('zulip')
+            with self.assertRaises(ValidationError):
+                check_subdomain_available('zulip', allow_reserved_subdomain=True)
+
+        # Reserved ones should only work with the flag
+        with self.assertRaises(ValidationError):
+            check_subdomain_available('stream')
+        check_subdomain_available('stream', allow_reserved_subdomain=True)
 
 class UserSignUpTest(InviteUserBase):