From 4c1496b4a4d380fa62ada223819ec3cd0d523bc8 Mon Sep 17 00:00:00 2001
From: Matthieu Rigal <matt.rigal@gmail.com>
Date: Fri, 12 Jun 2015 13:10:36 +0200
Subject: [PATCH] Updated URL and Email field regex validators, added schemes
 arg to urlfield

---
 mongoengine/fields.py  | 23 ++++++++++++++++-------
 tests/fields/fields.py | 18 +++++++++++++++++-
 2 files changed, 33 insertions(+), 8 deletions(-)

diff --git a/mongoengine/fields.py b/mongoengine/fields.py
index 9176828c..b685534d 100644
--- a/mongoengine/fields.py
+++ b/mongoengine/fields.py
@@ -119,22 +119,31 @@ class URLField(StringField):
     """
 
     _URL_REGEX = re.compile(
-        r'^(?:http|ftp)s?://'  # http:// or https://
-        # domain...
-        r'(?:(?:[A-Z0-9](?:[A-Z0-9-]{0,61}[A-Z0-9])?\.)+(?:[A-Z]{2,6}\.?|[A-Z0-9-]{2,}\.?)|'
+        r'^(?:[a-z0-9\.\-]*)://'  # scheme is validated separately
+        r'(?:(?:[A-Z0-9](?:[A-Z0-9-]{0,61}[A-Z0-9])?\.)+(?:[A-Z]{2,6}\.?|[A-Z0-9-]{2,}(?<!-)\.?)|'  # domain...
         r'localhost|'  # localhost...
-        r'\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3})'  # ...or ip
+        r'\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3}|'  # ...or ipv4
+        r'\[?[A-F0-9]*:[A-F0-9:]+\]?)'  # ...or ipv6
         r'(?::\d+)?'  # optional port
         r'(?:/?|[/?]\S+)$', re.IGNORECASE)
+    _URL_SCHEMES = ['http', 'https', 'ftp', 'ftps']
 
-    def __init__(self, verify_exists=False, url_regex=None, **kwargs):
+    def __init__(self, verify_exists=False, url_regex=None, schemes=None, **kwargs):
         self.verify_exists = verify_exists
         self.url_regex = url_regex or self._URL_REGEX
+        self.schemes = schemes or self._URL_SCHEMES
         super(URLField, self).__init__(**kwargs)
 
     def validate(self, value):
+        # Check first if the scheme is valid
+        scheme = value.split('://')[0].lower()
+        if scheme not in self.schemes:
+            self.error('Invalid scheme {} in URL: {}'.format(scheme, value))
+            return
+
+        # Then check full URL
         if not self.url_regex.match(value):
-            self.error('Invalid URL: %s' % value)
+            self.error('Invalid URL: {}'.format(value))
             return
 
         if self.verify_exists:
@@ -162,7 +171,7 @@ class EmailField(StringField):
         # quoted-string
         r'|^"([\001-\010\013\014\016-\037!#-\[\]-\177]|\\[\001-011\013\014\016-\177])*"'
         # domain (max length of an ICAAN TLD is 22 characters)
-        r')@(?:[A-Z0-9](?:[A-Z0-9-]{0,253}[A-Z0-9])?\.)+[A-Z]{2,22}$', re.IGNORECASE
+        r')@(?:[A-Z0-9](?:[A-Z0-9-]{0,61}[A-Z0-9])?\.)+(?:[A-Z]{2,6}|[A-Z0-9-]{2,}(?<!-))$', re.IGNORECASE
     )
 
     def validate(self, value):
diff --git a/tests/fields/fields.py b/tests/fields/fields.py
index fd083c73..7e04482f 100644
--- a/tests/fields/fields.py
+++ b/tests/fields/fields.py
@@ -341,6 +341,23 @@ class FieldTest(unittest.TestCase):
         link.url = 'http://www.google.com:8080'
         link.validate()
 
+    def test_url_scheme_validation(self):
+        """Ensure that URLFields validate urls with specific schemes properly.
+        """
+        class Link(Document):
+            url = URLField()
+
+        class SchemeLink(Document):
+            url = URLField(schemes=['ws', 'irc'])
+
+        link = Link()
+        link.url = 'ws://google.com'
+        self.assertRaises(ValidationError, link.validate)
+
+        scheme_link = SchemeLink()
+        scheme_link.url = 'ws://google.com'
+        scheme_link.validate()
+
     def test_int_validation(self):
         """Ensure that invalid values cannot be assigned to int fields.
         """
@@ -3088,7 +3105,6 @@ class FieldTest(unittest.TestCase):
         self.assertTrue(user.validate() is None)
 
         user = User(email=("Kofq@rhom0e4klgauOhpbpNdogawnyIKvQS0wk2mjqrgGQ5S"
-                           "ucictfqpdkK9iS1zeFw8sg7s7cwAF7suIfUfeyueLpfosjn3"
                            "aJIazqqWkm7.net"))
         self.assertTrue(user.validate() is None)