From 1e6baf57f6c56dd82bf2596e746682b8bd5b18cd Mon Sep 17 00:00:00 2001
From: tcsenpai <thecookingsenpai@countermail.com>
Date: Sat, 10 Feb 2024 22:40:45 +0100
Subject: [PATCH] Implemented AES256

---
 README.md                | 38 +++++++++++++++++++++++++++++++++++-
 src/hmacrypt.py          | 21 +++++++++++---------
 src/libs/seedable_aes.py | 42 ++++++++++++++++++++++++++++++++++++++++
 3 files changed, 91 insertions(+), 10 deletions(-)
 create mode 100644 src/libs/seedable_aes.py

diff --git a/README.md b/README.md
index 98941c0..5341d8f 100644
--- a/README.md
+++ b/README.md
@@ -1,5 +1,19 @@
 # hmacrypt
-encryption / decryption / signing / verifying using RSA and ECDSA a deterministic Key Derivation from password+hmac_secret by an hardware key.
+encryption / decryption / signing / verifying using RSA, AES and ECDSA a deterministic Key Derivation from password+hmac_secret by an hardware key.
+
+## Supported algorithms
+
+- RSA (deterministic with keyfile, password and hw device)
+- ECDSA (deterministic with keyfile, password and hw device)
+- AES 256 (deterministic with keyfile, password and hw device)
+
+## Tested devices
+
+- Thetis FIDO2
+
+### Should also work with
+
+Any FIDO2/U2F key having hmac_secrets extension
 
 ## Requirements (BEFORE everything else) & Credits
 
@@ -36,6 +50,7 @@ On Ubuntu 23.10 (untested on other platforms and flavors):
 
 - 2fa encryption/decryption using RSA deterministic Key Derivation
 - 2fa signing/verifying using ECDSA deterministic Key Derivation
+- 2fa AES cipher encryption/decryption using SHA256
 - Possibility to use the same or different keyfiles to enhance security
 - Consequently, supports for a wide range of hardware keys as long as they are compatible with libfido2
 - Low footprint: requires just two (or three) python library and a single system library
@@ -96,6 +111,27 @@ Or the path you used for the library.
 *If you REALLY have to change the src directory name, please correct the various paths inside.*
 
 
+#### encrypt_aes
+
+Definition:
+    
+    def encrypt_aes(message)
+
+Parameters:
+
+- message; data (string preferrable) to be encrypted
+
+#### decrypt_aes
+
+Definition:
+
+    def decrypt_aes(encrypted)
+
+Parameters:
+
+- encrypted; (usually) bytes to be decrypted
+
+
 #### inferECDSAKeys
 
 Definition:
diff --git a/src/hmacrypt.py b/src/hmacrypt.py
index 7aec5f6..6d57e44 100644
--- a/src/hmacrypt.py
+++ b/src/hmacrypt.py
@@ -1,6 +1,7 @@
 import subprocess
 from src.libs.seedable_rsa import decrypt, encrypt, generate_rsa_key
 from src.libs.seedable_ecdsa import generate_ecdsa_key, sign, verify
+from src.libs.seedable_aes import self_encrypt_aes, self_decrypt_aes
 
 # INFO This method derives the HMAC secret from the hardware key and the stored secret
 def getHMACSecret(keyfilePath="src/bins/.keyfile"):
@@ -15,6 +16,17 @@ def getHMACSecret(keyfilePath="src/bins/.keyfile"):
     hmac_secret = hmac_secret.strip()
     return hmac_secret
 
+# INFO The following methods are proxies to the AES methods
+# NOTE The AES methods generate a cipher on the fly based on the getHMACSecret method
+
+def encrypt_aes(message):
+    seed = getHMACSecret()
+    return self_encrypt_aes(seed, message)
+
+def decrypt_aes(encrypted):
+    seed = getHMACSecret()
+    return self_decrypt_aes(seed, encrypted)
+
 # INFO This method derives an ECDSA keypair from the stored secret and the hardware key
 def inferECDSAKeys(hidePrivate=False, savePublic=False):
     hmac_secret = getHMACSecret()
@@ -108,12 +120,3 @@ def self_decrypt_file(filepath, outpath):
 
 
 # TODO LARGE FILES
-
-
-# Self testing
-if __name__ == "__main__":
-    private_key, public_key = inferRSAKeys()
-    secret = encrypt("secret message", public_key)
-    print(secret)
-    decrypted = decrypt(secret, private_key)
-    print(decrypted)
diff --git a/src/libs/seedable_aes.py b/src/libs/seedable_aes.py
new file mode 100644
index 0000000..1ac8095
--- /dev/null
+++ b/src/libs/seedable_aes.py
@@ -0,0 +1,42 @@
+import base64
+import hashlib
+from Crypto import Random
+from Crypto.Cipher import AES
+
+class AESCipher(object):
+
+    def __init__(self, key):
+        self.bs = AES.block_size
+        self.key = hashlib.sha256(key.encode()).digest()
+
+    def encrypt(self, raw):
+        raw = self._pad(raw)
+        iv = Random.new().read(AES.block_size)
+        cipher = AES.new(self.key, AES.MODE_CBC, iv)
+        return base64.b64encode(iv + cipher.encrypt(raw.encode()))
+
+    def decrypt(self, enc):
+        enc = base64.b64decode(enc)
+        iv = enc[:AES.block_size]
+        cipher = AES.new(self.key, AES.MODE_CBC, iv)
+        return AESCipher._unpad(cipher.decrypt(enc[AES.block_size:])).decode('utf-8')
+
+    def _pad(self, s):
+        return s + (self.bs - len(s) % self.bs) * chr(self.bs - len(s) % self.bs)
+
+    @staticmethod
+    def _unpad(s):
+        return s[:-ord(s[len(s)-1:])]
+    
+# Implementable methods
+# NOTE The cipher is initialized on the fly, so the seed is not stored
+
+def self_encrypt_aes(seed, message):
+    cipher = AESCipher(seed)
+    encrypted = cipher.encrypt(message)
+    return encrypted
+
+def self_decrypt_aes(seed, encrypted):
+    cipher = AESCipher(seed)
+    decrypted = cipher.decrypt(encrypted)
+    return decrypted