add MGF1 class, docs, tests

3 files changed, 66 insertions, 0 deletions

diff --git a/cryptography/hazmat/primitives/asymmetric/padding.py b/cryptography/hazmat/primitives/asymmetric/padding.py
index 6bafe314..df960b0f 100644
--- a/cryptography/hazmat/primitives/asymmetric/padding.py
+++ b/cryptography/hazmat/primitives/asymmetric/padding.py
@@ -13,6 +13,8 @@
 
 from __future__ import absolute_import, division, print_function
 
+import six
+
 from cryptography import utils
 from cryptography.hazmat.primitives import interfaces
 
@@ -20,3 +22,22 @@ from cryptography.hazmat.primitives import interfaces
 @utils.register_interface(interfaces.AsymmetricPadding)
 class PKCS1v15(object):
     name = "EMSA-PKCS1-v1_5"
+
+
+class MGF1(object):
+    MAX_LENGTH = b"MAX_LENGTH"
+
+    def __init__(self, algorithm, salt_length):
+        if not isinstance(algorithm, interfaces.HashAlgorithm):
+            raise TypeError("Expected instance of interfaces.HashAlgorithm.")
+
+        self.algorithm = algorithm
+
+        if (not isinstance(salt_length, six.integer_types) and
+                salt_length != b"MAX_LENGTH"):
+            raise TypeError("salt_length must be an integer")
+
+        if salt_length != b"MAX_LENGTH" and salt_length < 0:
+            raise ValueError("salt_length must be zero or greater")
+
+        self.salt_length = salt_length
diff --git a/docs/hazmat/primitives/asymmetric/padding.rst b/docs/hazmat/primitives/asymmetric/padding.rst
index 7aec3bd3..632f6985 100644
--- a/docs/hazmat/primitives/asymmetric/padding.rst
+++ b/docs/hazmat/primitives/asymmetric/padding.rst
@@ -17,4 +17,22 @@ Padding
     PKCS1 v1.5 (also known as simply PKCS1) is a simple padding scheme
     developed for use with RSA keys. It is defined in :rfc:`3447`.
 
+Mask Generation Functions
+~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. class:: MGF1(algorithm, salt_length)
+
+    .. versionadded:: 0.3
+
+    MGF1 (Mask Generation Function 1) is used as the mask generation function
+    in :class:`PSS` padding. It takes a hash algorithm and a salt length.
+
+    :param algorithm: An instance of a
+        :class:`~cryptography.hazmat.primitives.interfaces.HashAlgorithm`
+        provider.
+
+    :param int salt_length: The length of the salt. It is recommended that this
+        be set to ``MGF1.MAX_LENGTH``.
+
+
 .. _`Padding is critical`: http://rdist.root.org/2009/10/06/why-rsa-encryption-padding-is-critical/
diff --git a/tests/hazmat/primitives/test_rsa.py b/tests/hazmat/primitives/test_rsa.py
index 79323265..c4955478 100644
--- a/tests/hazmat/primitives/test_rsa.py
+++ b/tests/hazmat/primitives/test_rsa.py
@@ -558,3 +558,30 @@ class TestRSAVerification(object):
         public_key = private_key.public_key()
         with pytest.raises(TypeError):
             public_key.verifier(b"sig", "notpadding", hashes.SHA1(), backend)
+
+
+class TestMGF1(object):
+    def test_invalid_hash_algorithm(self):
+        with pytest.raises(TypeError):
+            padding.MGF1(b"not_a_hash", 0)
+
+    def test_invalid_salt_length_not_integer(self):
+        with pytest.raises(TypeError):
+            padding.MGF1(hashes.SHA1(), b"not_a_length")
+
+    def test_invalid_salt_length_negative_integer(self):
+        with pytest.raises(ValueError):
+            padding.MGF1(hashes.SHA1(), -1)
+
+    def test_valid_mgf1_parameters(self):
+        algorithm = hashes.SHA1()
+        salt_length = algorithm.digest_size
+        mgf = padding.MGF1(algorithm, salt_length)
+        assert mgf.algorithm == algorithm
+        assert mgf.salt_length == salt_length
+
+    def test_valid_mgf1_parameters_maximum(self):
+        algorithm = hashes.SHA1()
+        mgf = padding.MGF1(algorithm, padding.MGF1.MAX_LENGTH)
+        assert mgf.algorithm == algorithm
+        assert mgf.salt_length == padding.MGF1.MAX_LENGTH