pss.py

import logging

from cryptography.exceptions import InvalidSignature
from cryptography.hazmat.backends import default_backend
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.asymmetric import padding
from cryptography.hazmat.primitives.asymmetric import utils

from ..exception import BadSignature
from ..exception import Unsupported
from . import Signer

logger = logging.getLogger(__name__)


class PSSSigner(Signer):
    def __init__(self, algorithm="SHA256"):
        if algorithm == "SHA256":
            self.hash_algorithm = hashes.SHA256
        elif algorithm == "SHA384":
            self.hash_algorithm = hashes.SHA384
        elif algorithm == "SHA512":
            self.hash_algorithm = hashes.SHA512
        else:
            raise Unsupported("algorithm: {}".format(algorithm))

    def sign(self, msg, key):
        """
        Create a signature over a message

        :param msg: The message
        :param key: The key
        :return: A signature
        """
        hasher = hashes.Hash(self.hash_algorithm(), backend=default_backend())
        hasher.update(msg)
        digest = hasher.finalize()
        sig = key.sign(
            digest,
            padding.PSS(
                mgf=padding.MGF1(self.hash_algorithm()),
                salt_length=padding.PSS.MAX_LENGTH,
            ),
            utils.Prehashed(self.hash_algorithm()),
        )
        return sig

    def verify(self, msg, signature, key):
        """
        Verify a message signature

        :param msg: The message
        :param sig: A signature
        :param key: A ec.EllipticCurvePublicKey to use for the verification.
        :raises: BadSignature if the signature can't be verified.
        :return: True
        """
        try:
            key.verify(
                signature,
                msg,
                padding.PSS(
                    mgf=padding.MGF1(self.hash_algorithm()),
                    salt_length=padding.PSS.MAX_LENGTH,
                ),
                self.hash_algorithm(),
            )
        except InvalidSignature as err:
            raise BadSignature(err)
        else:
            return True