bpo-31702: Allow to specify rounds for SHA-2 hashing in crypt.mksalt(). (#4110)

The log_rounds parameter for Blowfish has been replaced with the rounds parameter.

5 files changed, 74 insertions, 28 deletions

diff --git a/Doc/library/crypt.rst b/Doc/library/crypt.rst
index 9877b71..dd62cb3 100644
--- a/Doc/library/crypt.rst
+++ b/Doc/library/crypt.rst
@@ -116,7 +116,7 @@ The :mod:`crypt` module defines the following functions:
       Accept ``crypt.METHOD_*`` values in addition to strings for *salt*.
 
 
-.. function:: mksalt(method=None, *, log_rounds=12)
+.. function:: mksalt(method=None, *, rounds=None)
 
    Return a randomly generated salt of the specified method.  If no
    *method* is given, the strongest method available as returned by
@@ -125,14 +125,18 @@ The :mod:`crypt` module defines the following functions:
    The return value is a string suitable for passing as the *salt* argument
    to :func:`crypt`.
 
-   *log_rounds* specifies the binary logarithm of the number of rounds
-   for ``crypt.METHOD_BLOWFISH``, and is ignored otherwise.  ``8`` specifies
-   ``256`` rounds.
+   *rounds* specifies the number of rounds for ``METHOD_SHA256``,
+   ``METHOD_SHA512`` and ``METHOD_BLOWFISH``.
+   For ``METHOD_SHA256`` and ``METHOD_SHA512`` it must be an integer between
+   ``1000`` and ``999_999_999``, the default is ``5000``.  For
+   ``METHOD_BLOWFISH`` it must be a power of two between ``16`` (2\ :sup:`4`)
+   and ``2_147_483_648`` (2\ :sup:`31`), the default is ``4096``
+   (2\ :sup:`12`).
 
    .. versionadded:: 3.3
 
    .. versionchanged:: 3.7
-      Added the *log_rounds* parameter.
+      Added the *rounds* parameter.
 
 
 Examples
diff --git a/Doc/whatsnew/3.7.rst b/Doc/whatsnew/3.7.rst
index bb75939..71e8358 100644
--- a/Doc/whatsnew/3.7.rst
+++ b/Doc/whatsnew/3.7.rst
@@ -280,6 +280,9 @@ crypt
 Added support for the Blowfish method.
 (Contributed by Serhiy Storchaka in :issue:`31664`.)
 
+The :func:`~crypt.mksalt` function now allows to specify the number of rounds
+for hashing.  (Contributed by Serhiy Storchaka in :issue:`31702`.)
+
 dis
 ---
 
diff --git a/Lib/crypt.py b/Lib/crypt.py
index 4d73202..b0e47f4 100644
--- a/Lib/crypt.py
+++ b/Lib/crypt.py
@@ -19,7 +19,7 @@ class _Method(_namedtuple('_Method', 'name ident salt_chars total_size')):
         return '<crypt.METHOD_{}>'.format(self.name)
 
 
-def mksalt(method=None, *, log_rounds=12):
+def mksalt(method=None, *, rounds=None):
     """Generate a salt for the specified method.
 
     If not specified, the strongest available method will be used.
@@ -27,12 +27,32 @@ def mksalt(method=None, *, log_rounds=12):
     """
     if method is None:
         method = methods[0]
-    if not method.ident:
+    if rounds is not None and not isinstance(rounds, int):
+        raise TypeError(f'{rounds.__class__.__name__} object cannot be '
+                        f'interpreted as an integer')
+    if not method.ident:  # traditional
         s = ''
-    elif method.ident[0] == '2':
-        s = f'${method.ident}${log_rounds:02d}$'
-    else:
+    else:  # modular
         s = f'${method.ident}$'
+
+    if method.ident and method.ident[0] == '2':  # Blowfish variants
+        if rounds is None:
+            log_rounds = 12
+        else:
+            log_rounds = int.bit_length(rounds-1)
+            if rounds != 1 << log_rounds:
+                raise ValueError('rounds must be a power of 2')
+            if not 4 <= log_rounds <= 31:
+                raise ValueError('rounds out of the range 2**4 to 2**31')
+        s += f'{log_rounds:02d}$'
+    elif method.ident in ('5', '6'):  # SHA-2
+        if rounds is not None:
+            if not 1000 <= rounds <= 999_999_999:
+                raise ValueError('rounds out of the range 1000 to 999_999_999')
+            s += f'rounds={rounds}$'
+    elif rounds is not None:
+        raise ValueError(f"{method} doesn't support the rounds argument")
+
     s += ''.join(_sr.choice(_saltchars) for char in range(method.salt_chars))
     return s
 
@@ -55,10 +75,10 @@ def crypt(word, salt=None):
 #  available salting/crypto methods
 methods = []
 
-def _add_method(name, *args):
+def _add_method(name, *args, rounds=None):
     method = _Method(name, *args)
     globals()['METHOD_' + name] = method
-    salt = mksalt(method, log_rounds=4)
+    salt = mksalt(method, rounds=rounds)
     result = crypt('', salt)
     if result and len(result) == method.total_size:
         methods.append(method)
@@ -74,7 +94,7 @@ _add_method('SHA256', '5', 16, 63)
 # 'y' is the same as 'b', for compatibility
 # with openwall crypt_blowfish.
 for _v in 'b', 'y', 'a', '':
-    if _add_method('BLOWFISH', '2' + _v, 22, 59 + len(_v)):
+    if _add_method('BLOWFISH', '2' + _v, 22, 59 + len(_v), rounds=1<<4):
         break
 
 _add_method('MD5', '1', 8, 34)
diff --git a/Lib/test/test_crypt.py b/Lib/test/test_crypt.py
index 796fd07..d9189fc 100644
--- a/Lib/test/test_crypt.py
+++ b/Lib/test/test_crypt.py
@@ -39,12 +39,26 @@ class CryptTestCase(unittest.TestCase):
         else:
             self.assertEqual(crypt.methods[-1], crypt.METHOD_CRYPT)
 
+    @unittest.skipUnless(crypt.METHOD_SHA256 in crypt.methods or
+                         crypt.METHOD_SHA512 in crypt.methods,
+                        'requires support of SHA-2')
+    def test_sha2_rounds(self):
+        for method in (crypt.METHOD_SHA256, crypt.METHOD_SHA512):
+            for rounds in 1000, 10_000, 100_000:
+                salt = crypt.mksalt(method, rounds=rounds)
+                self.assertIn('$rounds=%d$' % rounds, salt)
+                self.assertEqual(len(salt) - method.salt_chars,
+                                 11 + len(str(rounds)))
+                cr = crypt.crypt('mypassword', salt)
+                self.assertTrue(cr)
+                cr2 = crypt.crypt('mypassword', cr)
+                self.assertEqual(cr2, cr)
+
     @unittest.skipUnless(crypt.METHOD_BLOWFISH in crypt.methods,
                         'requires support of Blowfish')
-    def test_log_rounds(self):
-        self.assertEqual(len(crypt._saltchars), 64)
+    def test_blowfish_rounds(self):
         for log_rounds in range(4, 11):
-            salt = crypt.mksalt(crypt.METHOD_BLOWFISH, log_rounds=log_rounds)
+            salt = crypt.mksalt(crypt.METHOD_BLOWFISH, rounds=1 << log_rounds)
             self.assertIn('$%02d$' % log_rounds, salt)
             self.assertIn(len(salt) - crypt.METHOD_BLOWFISH.salt_chars, {6, 7})
             cr = crypt.crypt('mypassword', salt)
@@ -52,18 +66,21 @@ class CryptTestCase(unittest.TestCase):
             cr2 = crypt.crypt('mypassword', cr)
             self.assertEqual(cr2, cr)
 
-    @unittest.skipUnless(crypt.METHOD_BLOWFISH in crypt.methods,
-                        'requires support of Blowfish')
-    def test_invalid_log_rounds(self):
-        for log_rounds in (1, -1, 999):
-            salt = crypt.mksalt(crypt.METHOD_BLOWFISH, log_rounds=log_rounds)
-            cr = crypt.crypt('mypassword', salt)
-            if cr is not None:
-                # On failure the openwall implementation returns a magic
-                # string that is shorter than 13 characters and is guaranteed
-                # to differ from a salt.
-                self.assertNotEqual(cr, salt)
-                self.assertLess(len(cr), 13)
+    def test_invalid_rounds(self):
+        for method in (crypt.METHOD_SHA256, crypt.METHOD_SHA512,
+                       crypt.METHOD_BLOWFISH):
+            with self.assertRaises(TypeError):
+                crypt.mksalt(method, rounds='4096')
+            with self.assertRaises(TypeError):
+                crypt.mksalt(method, rounds=4096.0)
+            for rounds in (0, 1, -1, 1<<999):
+                with self.assertRaises(ValueError):
+                    crypt.mksalt(method, rounds=rounds)
+        with self.assertRaises(ValueError):
+            crypt.mksalt(crypt.METHOD_BLOWFISH, rounds=1000)
+        for method in (crypt.METHOD_CRYPT, crypt.METHOD_MD5):
+            with self.assertRaisesRegex(ValueError, 'support'):
+                crypt.mksalt(method, rounds=4096)
 
 
 if __name__ == "__main__":
diff --git a/Misc/NEWS.d/next/Library/2017-10-24-21-10-44.bpo-31702.SfwJDI.rst b/Misc/NEWS.d/next/Library/2017-10-24-21-10-44.bpo-31702.SfwJDI.rst
new file mode 100644
index 0000000..3505cbd
--- /dev/null
+++ b/Misc/NEWS.d/next/Library/2017-10-24-21-10-44.bpo-31702.SfwJDI.rst
@@ -0,0 +1,2 @@
+crypt.mksalt() now allows to specify the number of rounds for SHA-256 and
+SHA-512 hashing.