基于Python rsa封装RSA加解密类

529 阅读4分钟

本文已参与「新人创作礼」活动,一起开启掘金创作之路。

前言

  • 本文为作者基于python rsa模块自行封装的RSA加解密工具代码
  • 支持以下功能:
    • 生成密钥
    • 保存密钥到文件
    • 从文件/代码加载密钥
    • 文本加密
    • 文本解密
    • base64编码解密

函数版

import rsa
import base64

'''
Author: 浅若清风cyf
Date: 2022/05/17
Description: 
    基于rsa模块封装了以下功能:
        生成密钥
        保存密钥到文件
        从文件/代码加载密钥
        文本加密
        文本解密
        base64编码解密
'''

'''
rsa模块:
    1、加解密:公钥加密、私钥解密
    2、签名验证:私钥加密,公钥验证
'''
pubBytes = b'''-----BEGIN RSA PUBLIC KEY-----
请将生成的公钥复制到此处
-----END RSA PUBLIC KEY-----'''

priBytes = b'''-----BEGIN RSA PRIVATE KEY-----
请将生成的私钥复制到此处
-----END RSA PRIVATE KEY-----'''


def gen_key(length=2048) -> (rsa.PublicKey, rsa.PrivateKey):
    publicKey, privateKey = rsa.newkeys(length)
    return publicKey, privateKey


def save_key(pubKey, priKey, pubFile='public_rsa.pem', priFile='private_rsa.pem'):
    try:
        if pubKey is not None:
            pub = pubKey.save_pkcs1()
            with open(pubFile, 'wb+')as f:
                f.write(pub)
        if priKey is not None:
            pri = priKey.save_pkcs1()
            with open(priFile, 'wb+')as f:
                f.write(pri)
        return True
    except Exception as e:
        print(e)
        return False


def load_key_from_file(pubFile=None, priFile=None) -> (rsa.PublicKey, rsa.PrivateKey):
    if pubFile is None and priFile is None:
        raise Exception('pubFile and priFile cannot be None at the same time.')

    if pubFile is not None:
        with open(pubFile, 'rb') as publicKeyFile:
            p = publicKeyFile.read()

        pubKey = rsa.PublicKey.load_pkcs1(p)

        if priFile is not None:
            with open(priFile, 'rb') as privateKeyFile:
                p = privateKeyFile.read()
            PriKey = rsa.PrivateKey.load_pkcs1(p)

            return pubKey, PriKey
        else:
            return pubKey, None
    else:
        with open(priFile, 'rb') as privateKeyFile:
            p = privateKeyFile.read()
        PriKey = rsa.PrivateKey.load_pkcs1(p)

        return None, PriKey


def load_key_from_code(pubBytes: bytes = None, priBytes: bytes = None) -> (rsa.PublicKey, rsa.PrivateKey):
    if pubBytes is None and priBytes is None:
        raise Exception('pubBytes and priBytes cannot be None at the same time.')

    if pubBytes is not None:
        pubKey = rsa.PublicKey.load_pkcs1(pubBytes)

        if priBytes is not None:
            PriKey = rsa.PrivateKey.load_pkcs1(priBytes)

            return pubKey, PriKey
        else:
            return pubKey, None
    else:
        PriKey = rsa.PrivateKey.load_pkcs1(priBytes)

        return None, PriKey


def encryption(text, pubKey: rsa.PublicKey, useBase64=False) -> bytes:
    # rsa加密,转base64字符串
    if isinstance(text, str):
        text = text.encode(encoding='utf-8')

    eText = rsa.encrypt(text, pubKey)

    if useBase64:
        eText = base64.b64encode(eText).decode(encoding='utf-8')

    return eText


def decryption(text, priKey: rsa.PrivateKey, useBase64=False):
    if useBase64:
        text.encode(encoding='utf-8')
        text = base64.b64decode(text)
    # base64字符串转bytes,rsa解密
    return rsa.decrypt(text, priKey).decode(encoding='utf-8')


def demo():
    publicKey, privateKey = rsa.newkeys(2048)
    message = "hello geeks"
    encMessage = rsa.encrypt(message.encode(encoding='utf-8'), publicKey)
    decMessage = rsa.decrypt(encMessage, privateKey).decode(encoding='utf-8')
    print(publicKey, privateKey)
    print(message)
    print(encMessage)
    print(decMessage)


if __name__ == '__main__':
    pubFile = 'public_rsa.pem'
    priFile = 'private_rsa.pem'
    # publicKey, privateKey = gen_key()
    # flag = save_key(publicKey, privateKey, pubFile, priFile)
    # print(flag)

    isBase64 = True
    isFromFile = False
    # isBase64 = False

    # no base64
    if not isBase64:
        if isFromFile:
            publicKey, priKey = load_key_from_file(pubFile, priFile)
        else:
            publicKey, priKey = load_key_from_code(pubBytes, priBytes)
        print(pubFile, priKey)
        enText = encryption('123456', publicKey)
        print(enText)
        deText = decryption(enText, priKey)
        print(deText)
    else:
        if isFromFile:
            publicKey, priKey = load_key_from_file(pubFile, priFile)
        else:
            publicKey, priKey = load_key_from_code(pubBytes, priBytes)
        print(pubFile, priKey)
        enText = encryption('123456', publicKey, useBase64=True)
        print(enText)
        deText = decryption(enText, priKey, useBase64=True)
        print(deText)

自定义类

import rsa
import base64

'''
Author: 浅若清风cyf
Date: 2022/05/17
Description: 
    基于rsa模块封装了以下功能:
        生成密钥
        保存密钥到文件
        从文件/代码加载密钥
        文本加密
        文本解密
        base64编码解密
'''

'''
rsa模块:
    1、加解密:公钥加密、私钥解密
    2、签名验证:私钥加密,公钥验证
'''
pubBytes = b'''-----BEGIN RSA PUBLIC KEY-----
请将生成的公钥复制到此处
-----END RSA PUBLIC KEY-----'''

priBytes = b'''-----BEGIN RSA PRIVATE KEY-----
请将生成的私钥复制到此处
-----END RSA PRIVATE KEY-----'''


class RSATools:

    def gen_key(self, length=2048) -> (rsa.PublicKey, rsa.PrivateKey):
        publicKey, privateKey = rsa.newkeys(length)
        return publicKey, privateKey

    def save_key(self, pubKey, priKey, pubFile='public_rsa.pem', priFile='private_rsa.pem'):
        try:
            if pubKey is not None:
                pub = pubKey.save_pkcs1()
                with open(pubFile, 'wb+')as f:
                    f.write(pub)
            if priKey is not None:
                pri = priKey.save_pkcs1()
                with open(priFile, 'wb+')as f:
                    f.write(pri)
            return True
        except Exception as e:
            print(e)
            return False

    def load_key_from_file(self, pubFile=None, priFile=None) -> (rsa.PublicKey, rsa.PrivateKey):
        if pubFile is None and priFile is None:
            raise Exception('pubFile and priFile cannot be None at the same time.')

        if pubFile is not None:
            with open(pubFile, 'rb') as publicKeyFile:
                p = publicKeyFile.read()

            pubKey = rsa.PublicKey.load_pkcs1(p)

            if priFile is not None:
                with open(priFile, 'rb') as privateKeyFile:
                    p = privateKeyFile.read()
                PriKey = rsa.PrivateKey.load_pkcs1(p)

                return pubKey, PriKey
            else:
                return pubKey, None
        else:
            with open(priFile, 'rb') as privateKeyFile:
                p = privateKeyFile.read()
            PriKey = rsa.PrivateKey.load_pkcs1(p)

            return None, PriKey

    def load_key_from_code(self, pubBytes: bytes = None, priBytes: bytes = None) -> (rsa.PublicKey, rsa.PrivateKey):
        if pubBytes is None and priBytes is None:
            raise Exception('pubBytes and priBytes cannot be None at the same time.')

        if pubBytes is not None:
            pubKey = rsa.PublicKey.load_pkcs1(pubBytes)

            if priBytes is not None:
                PriKey = rsa.PrivateKey.load_pkcs1(priBytes)

                return pubKey, PriKey
            else:
                return pubKey, None
        else:
            PriKey = rsa.PrivateKey.load_pkcs1(priBytes)

            return None, PriKey

    def encryption(self, text, pubKey: rsa.PublicKey, useBase64=False) -> bytes:
        # rsa加密,转base64字符串
        if isinstance(text, str):
            text = text.encode(encoding='utf-8')

        eText = rsa.encrypt(text, pubKey)

        if useBase64:
            eText = base64.b64encode(eText).decode(encoding='utf-8')

        return eText

    def decryption(self, text, priKey: rsa.PrivateKey, useBase64=False):
        if useBase64:
            text.encode(encoding='utf-8')
            text = base64.b64decode(text)
        # base64字符串转bytes,rsa解密
        return rsa.decrypt(text, priKey).decode(encoding='utf-8')


def demo():
    publicKey, privateKey = rsa.newkeys(2048)
    message = "hello geeks"
    encMessage = rsa.encrypt(message.encode(encoding='utf-8'), publicKey)
    decMessage = rsa.decrypt(encMessage, privateKey).decode(encoding='utf-8')
    print(publicKey, privateKey)
    print(message)
    print(encMessage)
    print(decMessage)


if __name__ == '__main__':
    pubFile = 'public_rsa.pem'
    priFile = 'private_rsa.pem'
    # publicKey, privateKey = gen_key()
    # flag = save_key(publicKey, privateKey, pubFile, priFile)
    # print(flag)

    isBase64 = True
    isFromFile = False
    # isBase64 = False
    rsaTools = RSATools()

    # no base64
    if not isBase64:
        if isFromFile:
            publicKey, priKey = rsaTools.load_key_from_file(pubFile, priFile)
        else:
            publicKey, priKey = rsaTools.load_key_from_code(pubBytes, priBytes)
        print(pubFile, priKey)
        enText = rsaTools.encryption('123456', publicKey)
        print(enText)
        deText = rsaTools.decryption(enText, priKey)
        print(deText)
    else:
        if isFromFile:
            publicKey, priKey = rsaTools.load_key_from_file(pubFile, priFile)
        else:
            publicKey, priKey = rsaTools.load_key_from_code(pubBytes, priBytes)
        print(pubFile, priKey)
        enText = rsaTools.encryption('123456', publicKey, useBase64=True)
        print(enText)
        deText = rsaTools.decryption(enText, priKey, useBase64=True)
        print(deText)