implement PASETO Version4 protocol

Author: purificant

paseto/protocol/version4.py

@@ -0,0 +1,209 @@
+"""
+This module contains Version4 implementation of the Paseto protocol.
+
+https://github.com/paseto-standard/paseto-spec/blob/master/docs/01-Protocol-Versions/Version4.md
+"""
+
+import hashlib
+import hmac
+import os
+from typing import Tuple
+
+from paseto.crypto import libsodium_wrapper, primitives
+from paseto.exceptions import InvalidKey, InvalidMac
+from paseto.paserk.keys import (
+    _TYPE_LOCAL,
+    _TYPE_PUBLIC,
+    _TYPE_SECRET,
+    _create_asymmetric_key,
+    _create_symmetric_key,
+    _deserialize_key,
+)
+from paseto.paserk.keys import _verify_key as _generic_verify_key
+from paseto.protocol.common import check_footer, check_header, decode_message
+from paseto.protocol.util import b64, pae
+
+HEADER_LOCAL = b"v4.local."
+HEADER_PUBLIC = b"v4.public."
+
+NONCE_SIZE = 32
+ENCRYPTION_KEY_LENGTH = 32
+AUTHENTICATION_KEY_LENGTH = 32
+MAC_SIZE = 32
+
+INFO_ENCRYPTION = b"paseto-encryption-key"
+INFO_AUTHENTICATION = b"paseto-auth-key-for-aead"
+
+
+def encrypt(
+    message: bytes, key: bytes, footer: bytes = b"", implicit_assertion: bytes = b""
+) -> bytes:
+    """PASETO Version4 encrypt function."""
+
+    # verify that key is intended for use with this function
+    _verify_key(key, _TYPE_LOCAL)
+    raw_key: bytes = _deserialize_key(key)
+
+    # Step 1
+    header: bytes = HEADER_LOCAL
+
+    # Step 2
+    nonce: bytes = os.urandom(NONCE_SIZE)
+
+    # Step 3
+    encryption_key: bytes
+    authentication_key: bytes
+    nonce2: bytes
+    encryption_key, authentication_key, nonce2 = _split_key(raw_key, nonce)
+
+    # Step 4
+    ciphertext: bytes = libsodium_wrapper.crypto_stream_xchacha20_xor(
+        message=message, nonce=nonce2, key=encryption_key
+    )
+
+    # Step 5
+    pre_auth: bytes = pae([header, nonce, ciphertext, footer, implicit_assertion])
+
+    # Step 6
+    message_authentication_code: bytes = hashlib.blake2b(
+        pre_auth, key=authentication_key, digest_size=32
+    ).digest()
+
+    # Step 7
+    ret: bytes = header + b64(nonce + ciphertext + message_authentication_code)
+    if footer:
+        ret += b"." + b64(footer)
+    return ret
+
+
+def decrypt(
+    message: bytes, key: bytes, footer: bytes = b"", implicit_assertion: bytes = b""
+) -> bytes:
+    """PASETO Version4 decrypt function."""
+
+    # verify that key is intended for use with this function
+    _verify_key(key, _TYPE_LOCAL)
+    raw_key: bytes = _deserialize_key(key)
+
+    # Step 1
+    check_footer(message, footer)
+
+    # Step 2
+    header: bytes = HEADER_LOCAL
+    check_header(message, header)
+
+    # Step 3
+    decoded: bytes = decode_message(message, len(header))
+    nonce: bytes = decoded[:NONCE_SIZE]
+    ciphertext: bytes = decoded[NONCE_SIZE:-MAC_SIZE]
+
+    # Step 4
+    encryption_key, authentication_key, nonce2 = _split_key(raw_key, nonce)
+
+    # Step 5
+    pre_auth: bytes = pae([header, nonce, ciphertext, footer, implicit_assertion])
+
+    # Step 6
+    computed_mac: bytes = hashlib.blake2b(
+        pre_auth, key=authentication_key, digest_size=MAC_SIZE
+    ).digest()
+
+    # Step 7
+    mac_in_message: bytes = decoded[-MAC_SIZE:]
+    if not hmac.compare_digest(mac_in_message, computed_mac):
+        raise InvalidMac("Invalid MAC for given ciphertext")
+
+    # Steps 8 and 9
+    return libsodium_wrapper.crypto_stream_xchacha20_xor(
+        message=ciphertext, nonce=nonce2, key=encryption_key
+    )
+
+
+def sign(
+    message: bytes,
+    secret_key: bytes,
+    footer: bytes = b"",
+    implicit_assertion: bytes = b"",
+) -> bytes:
+    """Sign message and return token which can then be used with verify()."""
+
+    # verify that key is intended for use with this function
+    _verify_key(secret_key, _TYPE_SECRET)
+    raw_secret_key: bytes = _deserialize_key(secret_key)
+
+    # Step 1
+    header = HEADER_PUBLIC
+
+    # Step 2
+    message2 = pae([header, message, footer, implicit_assertion])
+
+    # Step 3
+    signature = primitives.sign(message2, raw_secret_key)
+
+    # Step 4
+    ret = header + b64(message + signature)
+    if footer:
+        ret += b"." + b64(footer)
+
+    return ret
+
+
+def verify(
+    signed_message: bytes,
+    public_key: bytes,
+    footer: bytes = b"",
+    implicit_assertion: bytes = b"",
+) -> bytes:
+    """Verify signature and return message. Raises exception if signature is invalid."""
+
+    # verify that key is intended for use with this function
+    _verify_key(public_key, _TYPE_PUBLIC)
+    raw_public_key: bytes = _deserialize_key(public_key)
+
+    # Step 1
+    check_footer(signed_message, footer)
+
+    # Step 2
+    header = HEADER_PUBLIC
+    check_header(signed_message, header)
+
+    # Step 3
+    raw_inner_message: bytes = decode_message(signed_message, len(header))
+    signature = raw_inner_message[-64:]
+    message = raw_inner_message[:-64]
+
+    # Step 4
+    message2 = pae([header, message, footer, implicit_assertion])
+
+    # Steps 5 and 6
+    primitives.verify(signature, message2, raw_public_key)
+    return message
+
+
+def _split_key(key: bytes, nonce: bytes) -> Tuple[bytes, bytes, bytes]:
+
+    hashed: bytes = hashlib.blake2b(
+        INFO_ENCRYPTION + nonce, key=key, digest_size=56
+    ).digest()
+    encryption_key: bytes = hashed[:ENCRYPTION_KEY_LENGTH]
+    nonce2: bytes = hashed[ENCRYPTION_KEY_LENGTH:]
+    authentication_key: bytes = hashlib.blake2b(
+        INFO_AUTHENTICATION + nonce, key=key, digest_size=AUTHENTICATION_KEY_LENGTH
+    ).digest()
+
+    return encryption_key, authentication_key, nonce2
+
+
+def _verify_key(key: bytes, key_type: bytes) -> None:
+    if not _generic_verify_key(key, 4, key_type):
+        raise InvalidKey
+
+
+def create_symmetric_key() -> bytes:
+    """Return key for use with encrypt() and decrypt()."""
+    return _create_symmetric_key(4)
+
+
+def create_asymmetric_key() -> Tuple[bytes, bytes]:
+    """Return key pair for use with sign() and verify()."""
+    return _create_asymmetric_key(4)

tests/protocol/test_official_test_vectors_v4.py

@@ -0,0 +1,97 @@
+"""
+This module contains tests for official test vectors.
+
+Test vectors: https://github.com/paseto-standard/test-vectors
+Docs: https://github.com/paseto-standard/paseto-spec
+"""
+
+import os
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+from paseto.paserk.keys import _create_asymmetric_key, _create_symmetric_key
+from paseto.protocol import version4
+from tests.conftest import get_test_vector
+from tests.util import (
+    transform_test_case_for_v4_local,
+    transform_test_case_for_v4_public,
+)
+
+
+def get_test_cases(name: str):
+    """Return test cases filtered by name."""
+    return [
+        test_case
+        for test_case in get_test_vector("v4")["tests"]
+        if test_case["name"].startswith(name)
+    ]
+
+
+# use a test nonce for reproducible tests
+@patch.object(os, "urandom")
+# pylint: disable=too-many-arguments
+@pytest.mark.parametrize(
+    "test_name,nonce,raw_key_material,test_token,payload,footer,implicit_assertion",
+    [
+        transform_test_case_for_v4_local(test_case)
+        for test_case in get_test_cases("4-E")
+    ],
+)
+def test_v4_local(
+    mock: MagicMock,
+    test_name: str,
+    nonce: bytes,
+    raw_key_material: bytes,
+    test_token: bytes,
+    payload: bytes,
+    footer: bytes,
+    implicit_assertion: bytes,
+) -> None:
+    """Tests for v4.local (Shared-Key Encryption)."""
+
+    # use non random nonce for reproducible tests
+    mock.return_value = nonce
+    key = _create_symmetric_key(4, raw_key_material)
+
+    # verify that encrypt produces expected token
+    assert test_token == version4.encrypt(
+        payload, key, footer, implicit_assertion
+    ), test_name
+
+    # verify that decrypt produces expected payload
+    assert payload == version4.decrypt(
+        test_token, key, footer, implicit_assertion
+    ), test_name
+
+
+# pylint: disable=too-many-arguments
+@pytest.mark.parametrize(
+    "test_name,raw_public_key,raw_secret_key,test_token,payload,footer,implicit_assertion",
+    [
+        transform_test_case_for_v4_public(test_case)
+        for test_case in get_test_cases("4-S")
+    ],
+)
+def test_v4_public(
+    test_name: str,
+    raw_public_key: bytes,
+    raw_secret_key: bytes,
+    test_token: bytes,
+    payload: bytes,
+    footer: bytes,
+    implicit_assertion: bytes,
+) -> None:
+    """Tests for v4.public"""
+
+    public_key, secret_key = _create_asymmetric_key(4, raw_public_key, raw_secret_key)
+
+    # verify that sign produces expected token
+    assert test_token == version4.sign(
+        payload, secret_key, footer, implicit_assertion
+    ), test_name
+
+    # verify that token contains expected payload
+    assert payload == version4.verify(
+        test_token, public_key, footer, implicit_assertion
+    ), test_name

tests/protocol/test_version4.py

@@ -0,0 +1,49 @@
+"""This module contains test for version4.py"""
+
+import pytest
+
+from paseto.exceptions import InvalidKey, InvalidMac
+from paseto.paserk.keys import _create_symmetric_key
+from paseto.protocol import version4
+from paseto.protocol.version4 import _verify_key
+
+
+def test_encrypt_decrypt() -> None:
+    """Test that decrypt() reverses encrypt()."""
+    message: bytes = b"foo"
+    key: bytes = version4.create_symmetric_key()
+
+    token: bytes = version4.encrypt(message, key)
+    plain_text: bytes = version4.decrypt(token, key)
+
+    assert plain_text == message
+
+
+@pytest.mark.parametrize(
+    "footer,implicit_assertion", [(b"", b""), (b"some footer", b"some assertion")]
+)
+def test_sign_verify(footer: bytes, implicit_assertion: bytes) -> None:
+    """Check that verify() reverses sign()."""
+    public_key, secret_key = version4.create_asymmetric_key()
+    message = b"foo"
+
+    signed = version4.sign(message, secret_key, footer, implicit_assertion)
+    assert version4.verify(signed, public_key, footer, implicit_assertion) == message
+
+
+def test_decrypt_invalid_mac() -> None:
+    """Test that exception is raised when mac is not valid."""
+    message: bytes = b"foo"
+    key: bytes = _create_symmetric_key(4, b"0" * 32)
+
+    token: bytes = version4.encrypt(message, key)
+    # tamper with mac
+    token_with_invalid_mac = token[:40] + b"0" + token[41:]
+    with pytest.raises(InvalidMac):
+        version4.decrypt(token_with_invalid_mac, key)
+
+
+def test_verify_key() -> None:
+    """Test that exception is raised when key is not verified."""
+    with pytest.raises(InvalidKey):
+        _verify_key(b"", b"some type")