rsa.py

from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.asymmetric import padding

from . import Encrypter


class RSAEncrypter(Encrypter):
    def encrypt(self, msg, key, sign_padding="pkcs1_padding"):
        _chosen_hash = hashes.SHA1
        if sign_padding == "pkcs1_padding":
            _padding = padding.PKCS1v15
            return key.encrypt(msg, _padding())
        elif sign_padding == "pkcs1_oaep_padding":
            _padding = padding.OAEP
        elif sign_padding == "pkcs1_oaep_256_padding":
            _padding = padding.OAEP
            _chosen_hash = hashes.SHA256
        else:
            raise Exception("Unsupported padding")
        return key.encrypt(
            msg,
            _padding(
                mgf=padding.MGF1(algorithm=_chosen_hash()),
                algorithm=_chosen_hash(),
                label=None,
            ),
        )

    def decrypt(self, ciphertext, key, sign_padding="pkcs1_padding"):
        _chosen_hash = hashes.SHA1
        if sign_padding == "pkcs1_padding":
            _padding = padding.PKCS1v15
            return key.decrypt(ciphertext, _padding())
        elif sign_padding == "pkcs1_oaep_padding":
            _padding = padding.OAEP
        elif sign_padding == "pkcs1_oaep_256_padding":
            _padding = padding.OAEP
            _chosen_hash = hashes.SHA256
        else:
            raise Exception("Unsupported padding")

        try:
            text = key.decrypt(
                ciphertext,
                _padding(
                    mgf=padding.MGF1(algorithm=_chosen_hash()),
                    algorithm=_chosen_hash(),
                    label=None,
                ),
            )
        except Exception:
            raise

        return text