python 加密解密模块_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是我随便定义的一个网络。

torch的save方法是用二进制数据保存的torch模型，我们也只能用二进制方法加密和解密。

torch的load方法只能是load一个文件，并不能load一个对象，所以即便是你decrypte出来模型原文件，也只能先把它到文件里，然后在用torch load进来。

最后为了比较两个网络是否相同，我用了一个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)