python AES加密&解密 torch模型

工具

用于AES的crypto，二进制字符转换十六进制字符用的binascii，以及torch。
python 在Windows下使用AES时要安装的是pycryptodome 模块 

pip install pycryptodome 

python 在Linux下使用AES时要安装的是pycrypto模块  

pip install pycrypto

关于AES

crypto里面的AES加密需要一个初始向量v，在下面的代码里变量叫做iv，需要是byte类型；还需要加密/解密用的key，在代码里的变量叫key。

代码

Net是我随便定义的一个网络。

1. torch的save方法是用二进制数据保存的torch模型，我们也只能用二进制方法加密和解密。
2. torch的load方法只能是load一个文件，并不能load一个对象，所以即便是你decrypte出来模型原文件，也只能先把它到文件里，然后在用torch load进来。
3. 最后为了比较两个网络是否相同，我用了一个test_tensor去测试两个网络，得到相同的输出。注，比较前要先eval。具体见：这个博客

import torch
from Crypto.Cipher import AES
from binascii import b2a_hex, a2b_hex
from Util import Net

# 如果byte_string不足16位的倍数就用空格补足为16位
def add_to_16(byte_string):
    if len(byte_string) % 16:
        add = 16 - (len(byte_string) % 16)
    else:
        add = 0
    byte_string = byte_string + (b'\0' * add)
    return byte_string


key = '9999999999999999'.encode('utf-8')
mode = AES.MODE_CBC
iv = b'qqqqqqqqqqqqqqqq'



# 加密函数
def Encrypt(byte_string):
    byte_string = add_to_16(byte_string)
    cryptos = AES.new(key, mode, iv)
    cipher_text = cryptos.encrypt(byte_string)
    # 因为AES加密后的字符串不一定是ascii字符集的，输出保存可能存在问题，所以这里转为16进制字符串
    return b2a_hex(cipher_text)


# 解密后，去掉补足的空格用strip() 去掉
def Decrypt(byte_string):
    cryptos = AES.new(key, mode, iv)
    plain_text = cryptos.decrypt(a2b_hex(byte_string))
    return plain_text.rstrip(b'\0')


if __name__ == '__main__':
    #测试向量
    test_tensor=torch.randn(1,14)
    net=Net()
    #全部保存
    torch.save(net,'model')
    
    #加密&写加密文件    
    with open('model','rb') as f1:
        encrypted=Encrypt(f1.read())
        with open('model-encrypted','wb') as f2:
            f2.write(encrypted)
    #解密 加密过的文件
    with open('model-decrypted','wb') as f:
        content=open('model-encrypted','rb').read()
        f.write(Decrypt(content))
    
    #load原本model
    net1=torch.load('./model')
    net1=net1.eval()
    #load解密后的model
    net2=torch.load('model-decrypted')
    net2=net2.eval()
    #测试结果
    print(net2(test_tensor), net1(test_tensor))
    exit(0)