diff --git a/lite/pylite/test/test_global.py b/lite/pylite/test/test_global.py index d0bcfd4c..b30ae9bd 100644 --- a/lite/pylite/test/test_global.py +++ b/lite/pylite/test/test_global.py @@ -71,25 +71,3 @@ class TestGlobal(TestShuffleNet): network.load(model_path) self.do_forward(network) - - def test_update_decryption_key(self): - wrong_key = [0] * 32 - LiteGlobal.update_decryption_key("AES_default", wrong_key) - - with self.assertRaises(RuntimeError): - config = LiteConfig() - config.bare_model_cryption_name = "AES_default".encode("utf-8") - network = LiteNetwork(config) - model_path = os.path.join(self.source_dir, "shufflenet_crypt_aes.mge") - network.load(model_path) - - right_key = [i for i in range(32)] - LiteGlobal.update_decryption_key("AES_default", right_key) - - config = LiteConfig() - config.bare_model_cryption_name = "AES_default".encode("utf-8") - network = LiteNetwork(config) - model_path = os.path.join(self.source_dir, "shufflenet_crypt_aes.mge") - network.load(model_path) - - self.do_forward(network) diff --git a/lite/src/decryption/aes_decrypt.h b/lite/src/decryption/aes_decrypt.h index bc3af609..bdc33a50 100644 --- a/lite/src/decryption/aes_decrypt.h +++ b/lite/src/decryption/aes_decrypt.h @@ -41,10 +41,10 @@ public: } static std::vector get_decrypt_key() { - std::vector key = {0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, - 0x08, 0x09, 0x0A, 0x0B, 0x0C, 0x0D, 0x0E, 0x0F, - 0x10, 0x11, 0x12, 0x13, 0x14, 0x15, 0x16, 0x17, - 0x18, 0x19, 0x1A, 0x1B, 0x1C, 0x1D, 0x1E, 0x1F}; + std::vector key(32); + key = {0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09, 0x0A, + 0x0B, 0x0C, 0x0D, 0x0E, 0x0F, 0x10, 0x11, 0x12, 0x13, 0x14, 0x15, + 0x16, 0x17, 0x18, 0x19, 0x1A, 0x1B, 0x1C, 0x1D, 0x1E, 0x1F}; return key; } }; diff --git a/lite/tools/aes_encrypt.sh b/lite/tools/aes_encrypt.sh index 37900e1a..eabe89d9 100755 --- a/lite/tools/aes_encrypt.sh +++ b/lite/tools/aes_encrypt.sh @@ -11,6 +11,7 @@ fi IV=`openssl rand -hex 16` Key=000102030405060708090A0B0C0D0E0F101112131415161718191A1B1C1D1E1F + if [ $# == 3 ] ; then Key=$3 fi