From c4b785c29f32e54aa83df9743d2be06d5df910ae Mon Sep 17 00:00:00 2001
From: KillzXGaming <millernathan35@gmail.com>
Date: Wed, 17 Nov 2021 17:25:45 -0500
Subject: [PATCH] NKN : Fix alignment/padding for encrypting blocks fixing
 crashes.

---
 File_Format_Library/FileFormats/NKN.cs        | 16 ++++-
 .../Security/Cryptography/AesEncryption.cs    | 64 +++++++------------
 2 files changed, 37 insertions(+), 43 deletions(-)

diff --git a/File_Format_Library/FileFormats/NKN.cs b/File_Format_Library/FileFormats/NKN.cs
index 063b8654..0df8e5d0 100644
--- a/File_Format_Library/FileFormats/NKN.cs
+++ b/File_Format_Library/FileFormats/NKN.cs
@@ -85,12 +85,26 @@ namespace FirstPlugin
         public void Save(System.IO.Stream stream)
         {
             using (var writer = new BinaryWriter(stream)) {
-                writer.Write(AesEncryption.AesEncrypt(DecryptedContents));
+                AesEncryption.SetKey(AES_KEY);
+                AesEncryption.SetIV(AES_IV);
+                writer.Write(AesEncryption.AesEncrypt(IntoBytes(DecryptedContents)));
             }
         }
 
         public void Unload()
         {
         }
+
+        //Align the last set of bytes so everything gets encrypted back correctly
+        static byte[] IntoBytes(string contents)
+        {
+            var mem = new MemoryStream();
+            using (var writer = new FileWriter(mem))
+            {
+                writer.Write(Encoding.UTF8.GetBytes(contents));
+                writer.AlignBytes(128);
+            }
+            return mem.ToArray();
+        }
     }
 }
diff --git a/Switch_Toolbox_Library/Security/Cryptography/AesEncryption.cs b/Switch_Toolbox_Library/Security/Cryptography/AesEncryption.cs
index fe9f6392..9038ee11 100644
--- a/Switch_Toolbox_Library/Security/Cryptography/AesEncryption.cs
+++ b/Switch_Toolbox_Library/Security/Cryptography/AesEncryption.cs
@@ -18,11 +18,13 @@ namespace Toolbox.Library.Security.Cryptography
             IvBytes = UTF8Encoding.UTF8.GetBytes("000000000");
         }
 
-        public static void SetKey(string key) {
+        public static void SetKey(string key)
+        {
             keyBytes = UTF8Encoding.UTF8.GetBytes(key);
         }
 
-        public static void SetIV(string key) {
+        public static void SetIV(string key)
+        {
             IvBytes = UTF8Encoding.UTF8.GetBytes(key);
         }
 
@@ -52,53 +54,31 @@ namespace Toolbox.Library.Security.Cryptography
 
         public static string AesDecrypt(Byte[] inputBytes)
         {
-            Byte[] outputBytes = inputBytes;
-
-            string plaintext = string.Empty;
-
-            using (MemoryStream memoryStream = new MemoryStream(outputBytes))
-            {
-                using (CryptoStream cryptoStream = new CryptoStream(memoryStream, GetCryptoAlgorithm().CreateDecryptor(keyBytes, IvBytes), CryptoStreamMode.Read))
-                {
-                    using (StreamReader srDecrypt = new StreamReader(cryptoStream))
-                    {
-                        plaintext = srDecrypt.ReadToEnd();
-                    }
-                }
-            }
-
-            return plaintext;
+            var cryptoTransform = new AesManaged().CreateDecryptor(keyBytes, IvBytes);
+            return Encoding.UTF8.GetString(TransformBlocks(cryptoTransform, inputBytes));
         }
 
-        public static byte[] AesEncrypt(string inputText)
+        public static byte[] AesEncrypt(string inputText) => AesEncrypt(UTF8Encoding.UTF8.GetBytes(inputText));
+
+        public static byte[] AesEncrypt(byte[] inputBytes)
         {
-            byte[] inputBytes = UTF8Encoding.UTF8.GetBytes(inputText);//AbHLlc5uLone0D1q
+            var cryptoTransform = new AesManaged().CreateEncryptor(keyBytes, IvBytes);
+            return TransformBlocks(cryptoTransform, inputBytes);
+        }
 
-            byte[] result = null;
-            using (MemoryStream memoryStream = new MemoryStream())
+        static byte[] TransformBlocks(ICryptoTransform cryptoTransform, byte[] input)
+        {
+            byte[] result = new byte[input.Length];
+
+            int num = 0;
+            while (num < input.Length)
             {
-                using (CryptoStream cryptoStream = new CryptoStream(memoryStream, GetCryptoAlgorithm().CreateEncryptor(keyBytes, IvBytes), CryptoStreamMode.Write))
-                {
-                    cryptoStream.Write(inputBytes, 0, inputBytes.Length);
-                    cryptoStream.FlushFinalBlock();
-
-                    result = memoryStream.ToArray();
-                }
+                cryptoTransform.TransformBlock(input, num, 16, result, num);
+                num += 16;
             }
-
+            while (result[0] == (byte)0)
+                result = ((IEnumerable<byte>)result).Skip<byte>(1).ToArray<byte>();
             return result;
         }
-
-
-        private static RijndaelManaged GetCryptoAlgorithm()
-        {
-            RijndaelManaged algorithm = new RijndaelManaged();
-            //set the mode, padding and block size
-            algorithm.Padding = PaddingMode.PKCS7;
-            algorithm.Mode = CipherMode.CBC;
-            algorithm.KeySize = 128;
-            algorithm.BlockSize = 128;
-            return algorithm;
-        }
     }
 }