from Crypto.Cipher import AES
from binascii import b2a_hex, a2b_hex
# 如果text不足16位的倍數就用空格補足為16位
def add_to_16(text:bytes):
    if len(text) % 16:
        add = 16 - (len(text) % 16)
    else:
        add = 0
    text = text + (b'\0' * add)
    return text
# 加密函數
def encrypt(text:bytes):
    # key = '9999999999999999'.encode('utf-8')
    key = 'tqsktqsktqsktqsk'.encode('utf-8')
    mode = AES.MODE_CBC
    iv = b'qqqqqqqqqqqqqqqq'
    text = add_to_16(text)
    cryptos = AES.new(key, mode, iv)
    cipher_text = cryptos.encrypt(text)
    # 因為AES加密后的字符串不一定是ascii字符集的，輸出保存可能存在問題，所以這里轉為16進制字符串
    return b2a_hex(cipher_text)
# 解密后，去掉補足的空格用strip() 去掉
def decrypt(text):
    key = 'tqsktqsktqsktqsk'.encode('utf-8')
    iv = b'qqqqqqqqqqqqqqqq'
    mode = AES.MODE_CBC
    cryptos = AES.new(key, mode, iv)
    plain_text = cryptos.decrypt(a2b_hex(text))
    return plain_text
if __name__ == '__main__':
    # encryption model
    with open('detect_model/checkpoints/xxx_ori.pth', 'rb') as f1:
        encrypted = encrypt(f1.read())
        with open('detect_model/checkpoints/xxx_encryp.pth', 'wb') as f2:
            f2.write(encrypted)
    # decryption model
    with open("./detect_model/checkpoints/xxx_decryp.pth", 'wb') as f:
        content = open('detect_model/checkpoints/xxx_encryp.pth', 'rb').read()
        f.write(decrypt(content))