From fb3e02dcb558c230aef39251143d33b2a8b7b317 Mon Sep 17 00:00:00 2001
From: thezero <io@thezero.org>
Date: Sun, 19 May 2019 23:55:03 +0200
Subject: [PATCH] add AES padding options

---
 src/core/operations/AESDecrypt.mjs | 34 +++++++++++++++++++++++++-----
 src/core/operations/AESEncrypt.mjs | 19 +++++++++++++----
 2 files changed, 44 insertions(+), 9 deletions(-)

diff --git a/src/core/operations/AESDecrypt.mjs b/src/core/operations/AESDecrypt.mjs
index 53489d3c..6e979292 100644
--- a/src/core/operations/AESDecrypt.mjs
+++ b/src/core/operations/AESDecrypt.mjs
@@ -44,6 +44,11 @@ class AESDecrypt extends Operation {
                 "type": "option",
                 "value": ["CBC", "CFB", "OFB", "CTR", "GCM", "ECB"]
             },
+            {
+                "name": "Padding",
+                "type": "option",
+                "value": ["PKCS#7", "Null byte", "No padding"]
+            },
             {
                 "name": "Input",
                 "type": "option",
@@ -73,10 +78,8 @@ class AESDecrypt extends Operation {
     run(input, args) {
         const key = Utils.convertToByteArray(args[0].string, args[0].option),
             iv = Utils.convertToByteArray(args[1].string, args[1].option),
-            mode = args[2],
-            inputType = args[3],
-            outputType = args[4],
-            gcmTag = Utils.convertToByteString(args[5].string, args[5].option);
+            gcmTag = Utils.convertToByteString(args[6].string, args[6].option),
+            [,, mode, padding, inputType, outputType,] = args;
 
         if ([16, 24, 32].indexOf(key.length) < 0) {
             throw new OperationError(`Invalid key length: ${key.length} bytes
@@ -95,7 +98,28 @@ The following algorithms will be used based on the size of the key:
             tag: gcmTag
         });
         decipher.update(forge.util.createBuffer(input));
-        const result = decipher.finish();
+        var result = null;
+        if (padding === "PKCS#7") {
+            result = decipher.finish();
+        } else if (padding === "Null byte") {
+            result = decipher.finish(function(blockSize, buffer, decrypt) {
+                if (decrypt) {
+                    var len = buffer.length(), count = 0;
+                    for(var i = len - 1; i >= 8; --i) {
+                        if (buffer.at(i) == "00") {
+                            count += 1;
+                        } else {
+                            break;
+                        }
+                    }
+                    return buffer.truncate(count);
+                }
+            });
+        } else {
+            result = decipher.finish(function(blockSize, buffer, decrypt) {
+                return true;
+            });
+        };
 
         if (result) {
             return outputType === "Hex" ? decipher.output.toHex() : decipher.output.getBytes();
diff --git a/src/core/operations/AESEncrypt.mjs b/src/core/operations/AESEncrypt.mjs
index 5103cb86..e6c4eeca 100644
--- a/src/core/operations/AESEncrypt.mjs
+++ b/src/core/operations/AESEncrypt.mjs
@@ -44,6 +44,11 @@ class AESEncrypt extends Operation {
                 "type": "option",
                 "value": ["CBC", "CFB", "OFB", "CTR", "GCM", "ECB"]
             },
+            {
+                "name": "Padding",
+                "type": "option",
+                "value": ["PKCS#7", "Null byte"]
+            },
             {
                 "name": "Input",
                 "type": "option",
@@ -67,9 +72,7 @@ class AESEncrypt extends Operation {
     run(input, args) {
         const key = Utils.convertToByteArray(args[0].string, args[0].option),
             iv = Utils.convertToByteArray(args[1].string, args[1].option),
-            mode = args[2],
-            inputType = args[3],
-            outputType = args[4];
+            [,, mode, padding, inputType, outputType] = args;
 
         if ([16, 24, 32].indexOf(key.length) < 0) {
             throw new OperationError(`Invalid key length: ${key.length} bytes
@@ -85,7 +88,15 @@ The following algorithms will be used based on the size of the key:
         const cipher = forge.cipher.createCipher("AES-" + mode, key);
         cipher.start({iv: iv});
         cipher.update(forge.util.createBuffer(input));
-        cipher.finish();
+        if (padding === "PKCS#7") {
+            cipher.finish();
+        } else if (padding === "Null byte") {
+            cipher.finish(function(blockSize, buffer, decrypt) {
+                if (!decrypt) {
+                    return buffer.fillWithByte(0, blockSize);
+                }
+            });
+        }
 
         if (outputType === "Hex") {
             if (mode === "GCM") {