Merge pull request #192 from alex/pkcs7-padding

PKCS7 padding

6 files changed, 319 insertions, 8 deletions

diff --git a/cryptography/hazmat/primitives/interfaces.py b/cryptography/hazmat/primitives/interfaces.py
index 49c19d0e..217490fd 100644
--- a/cryptography/hazmat/primitives/interfaces.py
+++ b/cryptography/hazmat/primitives/interfaces.py
@@ -45,3 +45,17 @@ class CipherContext(six.with_metaclass(abc.ABCMeta)):
         """
         finalize return bytes
         """
+
+
+class PaddingContext(six.with_metaclass(abc.ABCMeta)):
+    @abc.abstractmethod
+    def update(self, data):
+        """
+        update takes bytes and return bytes
+        """
+
+    @abc.abstractmethod
+    def finalize(self):
+        """
+        finalize return bytes
+        """
diff --git a/cryptography/hazmat/primitives/padding.py b/cryptography/hazmat/primitives/padding.py
new file mode 100644
index 00000000..ddcadd89
--- /dev/null
+++ b/cryptography/hazmat/primitives/padding.py
@@ -0,0 +1,121 @@
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+# implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import six
+
+from cryptography.hazmat.primitives import interfaces
+
+
+class PKCS7(object):
+    def __init__(self, block_size):
+        super(PKCS7, self).__init__()
+        if not (0 <= block_size < 256):
+            raise ValueError("block_size must be in range(0, 256)")
+
+        if block_size % 8 != 0:
+            raise ValueError("block_size must be a multiple of 8")
+
+        self.block_size = block_size
+
+    def padder(self):
+        return _PKCS7PaddingContext(self.block_size)
+
+    def unpadder(self):
+        return _PKCS7UnpaddingContext(self.block_size)
+
+
+@interfaces.register(interfaces.PaddingContext)
+class _PKCS7PaddingContext(object):
+    def __init__(self, block_size):
+        super(_PKCS7PaddingContext, self).__init__()
+        self.block_size = block_size
+        # TODO: O(n ** 2) complexity for repeated concatentation, we should use
+        # zero-buffer (#193)
+        self._buffer = b""
+
+    def update(self, data):
+        if self._buffer is None:
+            raise ValueError("Context was already finalized")
+
+        if isinstance(data, six.text_type):
+            raise TypeError("Unicode-objects must be encoded before padding")
+
+        self._buffer += data
+
+        finished_blocks = len(self._buffer) // (self.block_size // 8)
+
+        result = self._buffer[:finished_blocks * (self.block_size // 8)]
+        self._buffer = self._buffer[finished_blocks * (self.block_size // 8):]
+
+        return result
+
+    def finalize(self):
+        if self._buffer is None:
+            raise ValueError("Context was already finalized")
+
+        pad_size = self.block_size // 8 - len(self._buffer)
+        result = self._buffer + six.int2byte(pad_size) * pad_size
+        self._buffer = None
+        return result
+
+
+@interfaces.register(interfaces.PaddingContext)
+class _PKCS7UnpaddingContext(object):
+    def __init__(self, block_size):
+        super(_PKCS7UnpaddingContext, self).__init__()
+        self.block_size = block_size
+        # TODO: O(n ** 2) complexity for repeated concatentation, we should use
+        # zero-buffer (#193)
+        self._buffer = b""
+
+    def update(self, data):
+        if self._buffer is None:
+            raise ValueError("Context was already finalized")
+
+        if isinstance(data, six.text_type):
+            raise TypeError("Unicode-objects must be encoded before unpadding")
+
+        self._buffer += data
+
+        finished_blocks = max(
+            len(self._buffer) // (self.block_size // 8) - 1,
+            0
+        )
+
+        result = self._buffer[:finished_blocks * (self.block_size // 8)]
+        self._buffer = self._buffer[finished_blocks * (self.block_size // 8):]
+
+        return result
+
+    def finalize(self):
+        if self._buffer is None:
+            raise ValueError("Context was already finalized")
+
+        if not self._buffer:
+            raise ValueError("Invalid padding bytes")
+
+        pad_size = six.indexbytes(self._buffer, -1)
+
+        if pad_size > self.block_size // 8:
+            raise ValueError("Invalid padding bytes")
+
+        mismatch = 0
+        for b in six.iterbytes(self._buffer[-pad_size:]):
+            mismatch |= b ^ pad_size
+
+        if mismatch != 0:
+            raise ValueError("Invalid padding bytes")
+
+        res = self._buffer[:-pad_size]
+        self._buffer = None
+        return res
diff --git a/docs/hazmat/primitives/index.rst b/docs/hazmat/primitives/index.rst
index 3927f3f0..ee1e251c 100644
--- a/docs/hazmat/primitives/index.rst
+++ b/docs/hazmat/primitives/index.rst
@@ -14,3 +14,4 @@ Primitives
     cryptographic-hashes
     hmac
     symmetric-encryption
+    padding
diff --git a/docs/hazmat/primitives/padding.rst b/docs/hazmat/primitives/padding.rst
new file mode 100644
index 00000000..ba3ddcc0
--- /dev/null
+++ b/docs/hazmat/primitives/padding.rst
@@ -0,0 +1,69 @@
+.. danger::
+
+    This is a "Hazardous Materials" module. You should **ONLY** use it if
+    you're 100% absolutely sure that you know what you're doing because this
+    module is full of land mines, dragons, and dinosaurs with laser guns.
+
+
+Padding
+=======
+
+.. currentmodule:: cryptography.hazmat.primitives.padding
+
+Padding is a way to take data that may or may not be be a multiple of the block
+size for a cipher and extend it out so that it is. This is required for many
+block cipher modes as they require the data to be encrypted to be an exact
+multiple of the block size.
+
+
+.. class:: PKCS7(block_size)
+
+    PKCS7 padding is a generalization of PKCS5 padding (also known as standard
+    padding). PKCS7 padding works by appending ``N`` bytes with the value of
+    ``chr(N)``, where ``N`` is the number of bytes required to make the final
+    block of data the same size as the block size. A simple example of padding
+    is:
+
+    .. doctest::
+
+        >>> from cryptography.hazmat.primitives import padding
+        >>> padder = padding.PKCS7(128).padder()
+        >>> padder.update(b"1111111111")
+        ''
+        >>> padder.finalize()
+        '1111111111\x06\x06\x06\x06\x06\x06'
+
+    :param block_size: The size of the block in bits that the data is being
+                       padded to.
+
+    .. method:: padder()
+
+        :returns: A padding
+            :class:`~cryptography.hazmat.primitives.interfaces.PaddingContext`
+            provider.
+
+    .. method:: unpadder()
+
+        :returns: An unpadding
+            :class:`~cryptography.hazmat.primitives.interfaces.PaddingContext`
+            provider.
+
+
+.. currentmodule:: cryptography.hazmat.primitives.interfaces
+
+.. class:: PaddingContext
+
+    When calling ``padder()`` or ``unpadder()`` you will receive an a return
+    object conforming to the ``PaddingContext`` interface. You can then call
+    ``update(data)`` with data until you have fed everything into the context.
+    Once that is done call ``finalize()`` to finish the operation and obtain
+    the remainder of the data.
+
+    .. method:: update(data)
+
+        :param bytes data: The data you wish to pass into the context.
+        :return bytes: Returns the data that was padded or unpadded.
+
+    .. method:: finalize()
+
+        :return bytes: Returns the remainder of the data.
diff --git a/docs/hazmat/primitives/symmetric-encryption.rst b/docs/hazmat/primitives/symmetric-encryption.rst
index 758a4648..9a5bce07 100644
--- a/docs/hazmat/primitives/symmetric-encryption.rst
+++ b/docs/hazmat/primitives/symmetric-encryption.rst
@@ -54,18 +54,17 @@ where the encrypter and decrypter both use the same key.
 
 .. currentmodule:: cryptography.hazmat.primitives.interfaces
 
-.. class:: CipherContext()
-
-    When calling ``encryptor()`` or ``decryptor()`` on a BlockCipher object you
-    will receive a return object conforming to the CipherContext interface. You
-    can then call ``update(data)`` with data until you have fed everything into
-    the context. Once that is done call ``finalize()`` to finish the operation and
-    obtain the remainder of the data.
+.. class:: CipherContext
 
+    When calling ``encryptor()`` or ``decryptor()`` on a ``BlockCipher`` object
+    you will receive a return object conforming to the ``CipherContext``
+    interface. You can then call ``update(data)`` with data until you have fed
+    everything into the context. Once that is done call ``finalize()`` to
+    finish the operation and obtain the remainder of the data.
 
     .. method:: update(data)
 
-        :param bytes data: The text you wish to pass into the context.
+        :param bytes data: The data you wish to pass into the context.
         :return bytes: Returns the data that was encrypted or decrypted.
 
     .. method:: finalize()
diff --git a/tests/hazmat/primitives/test_padding.py b/tests/hazmat/primitives/test_padding.py
new file mode 100644
index 00000000..3cefafaf
--- /dev/null
+++ b/tests/hazmat/primitives/test_padding.py
@@ -0,0 +1,107 @@
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+# implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import pytest
+
+import six
+
+from cryptography.hazmat.primitives import padding
+
+
+class TestPKCS7(object):
+    @pytest.mark.parametrize("size", [127, 4096, -2])
+    def test_invalid_block_size(self, size):
+        with pytest.raises(ValueError):
+            padding.PKCS7(size)
+
+    @pytest.mark.parametrize(("size", "padded"), [
+        (128, b"1111"),
+        (128, b"1111111111111111"),
+        (128, b"111111111111111\x06"),
+        (128, b""),
+    ])
+    def test_invalid_padding(self, size, padded):
+        unpadder = padding.PKCS7(size).unpadder()
+        with pytest.raises(ValueError):
+            unpadder.update(padded)
+            unpadder.finalize()
+
+    def test_non_bytes(self):
+        padder = padding.PKCS7(128).padder()
+        with pytest.raises(TypeError):
+            padder.update(six.u("abc"))
+        unpadder = padding.PKCS7(128).unpadder()
+        with pytest.raises(TypeError):
+            unpadder.update(six.u("abc"))
+
+    @pytest.mark.parametrize(("size", "unpadded", "padded"), [
+        (
+            128,
+            b"1111111111",
+            b"1111111111\x06\x06\x06\x06\x06\x06",
+        ),
+        (
+            128,
+            b"111111111111111122222222222222",
+            b"111111111111111122222222222222\x02\x02",
+        ),
+        (
+            128,
+            b"1" * 16,
+            b"1" * 16 + b"\x10" * 16,
+        ),
+        (
+            128,
+            b"1" * 17,
+            b"1" * 17 + b"\x0F" * 15,
+        )
+    ])
+    def test_pad(self, size, unpadded, padded):
+        padder = padding.PKCS7(size).padder()
+        result = padder.update(unpadded)
+        result += padder.finalize()
+        assert result == padded
+
+    @pytest.mark.parametrize(("size", "unpadded", "padded"), [
+        (
+            128,
+            b"1111111111",
+            b"1111111111\x06\x06\x06\x06\x06\x06",
+        ),
+        (
+            128,
+            b"111111111111111122222222222222",
+            b"111111111111111122222222222222\x02\x02",
+        ),
+    ])
+    def test_unpad(self, size, unpadded, padded):
+        unpadder = padding.PKCS7(size).unpadder()
+        result = unpadder.update(padded)
+        result += unpadder.finalize()
+        assert result == unpadded
+
+    def test_use_after_finalize(self):
+        padder = padding.PKCS7(128).padder()
+        b = padder.finalize()
+        with pytest.raises(ValueError):
+            padder.update(b"")
+        with pytest.raises(ValueError):
+            padder.finalize()
+
+        unpadder = padding.PKCS7(128).unpadder()
+        unpadder.update(b)
+        unpadder.finalize()
+        with pytest.raises(ValueError):
+            unpadder.update(b"")
+        with pytest.raises(ValueError):
+            unpadder.finalize()