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)