rsa_utils.py

#!/usr/bin/env python
# -*- encoding: utf-8 -*-
"""
@Version: 1.0
@Python Version:3.6.6
@Author: ludq1
@Email: ludq1@chinaunicom.cn
@date: 2023/04/07 11:40:00
@Description:
"""

import base64

import rsa


class CannotDecryptContentException(Exception):
    """
    自定义的异常类, 用于表示不能将加密后的内容进行界面的异常
    """

    def __init__(self, message: str):
        r"""
        初始化

        :param message
        """
        self.message: str = message

    def __str__(self):
        r"""
        字符串表示

        Returns:

        """
        return self.message

    def __repr__(self):
        r"""
        字符串表示

        Returns:

        """
        return self.message


class RsaUtils:
    r"""
    RSA utils类，不依赖任何其他应用基础类
    """

    def __init__(self, key: str = None):
        r"""
        初始化

        :param key: pkcs1格式，即 -----BEGIN RSA PRIVATE KEY-----
        """
        self.key: str = key

    def encrypt(self, to_encrypt_content: str, public_key: str):
        r"""
        RSA加密

        :param to_encrypt_content:  待加密的字符串
        :param public_key: 公钥，pkcs1格式,即 -----BEGIN RSA PUBLIC KEY-----
        :return: 加密后的字符串
        """
        to_encrypt_content = to_encrypt_content.encode("utf-8")
        rsa_public_key = rsa.PublicKey.load_pkcs1(public_key.encode("utf-8"), format='PEM')
        return base64.b64encode(rsa.encrypt(to_encrypt_content, rsa_public_key)).decode("utf-8")

    def gen_rsa_key(self) -> tuple:
        r"""
        生成公钥私钥

        :return:
        """
        return rsa.newkeys(512)

    def decrypt(self, to_decrypt_content: str, key: str = None) -> str:
        r"""
        RSA解密

        :param to_decrypt_content
        :param key：pkcs1格式，即 -----BEGIN RSA PRIVATE KEY-----

        """
        real_key = key if key else self.key
        private_key = rsa.PrivateKey.load_pkcs1(real_key.encode("utf-8"), format='PEM')
        try:
            content = rsa.decrypt(base64.b64decode(to_decrypt_content), private_key)
            return content.decode("utf-8")
        except BaseException as e:
            encountered_e = CannotDecryptContentException(str(e))
        if encountered_e:
            raise encountered_e