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.

1 files changed, 28 insertions, 8 deletions

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)