From 405b0e6fc5cf75fd85ee722d23e91203de7767ba Mon Sep 17 00:00:00 2001
From: TiagoRG <35657250+TiagoRG@users.noreply.github.com>
Date: Tue, 28 Feb 2023 18:36:12 +0000
Subject: [PATCH] AES cipher done (working for now at least)

---
 1ano/2semestre/labi/tema01/src/aes_cipher.py  | 42 +++++++++++++++
 1ano/2semestre/labi/tema01/src/aes_cypher.py  | 12 -----
 .../2semestre/labi/tema01/src/aes_decipher.py | 51 +++++++++++++++++++
 3 files changed, 93 insertions(+), 12 deletions(-)
 create mode 100644 1ano/2semestre/labi/tema01/src/aes_cipher.py
 delete mode 100644 1ano/2semestre/labi/tema01/src/aes_cypher.py
 create mode 100644 1ano/2semestre/labi/tema01/src/aes_decipher.py

diff --git a/1ano/2semestre/labi/tema01/src/aes_cipher.py b/1ano/2semestre/labi/tema01/src/aes_cipher.py
new file mode 100644
index 0000000..66f3ce6
--- /dev/null
+++ b/1ano/2semestre/labi/tema01/src/aes_cipher.py
@@ -0,0 +1,42 @@
+import os
+import sys
+from Crypto.Cipher import AES
+from Crypto.Hash import SHA256
+
+
+def main(args=None):
+    if len(args) < 2:
+        print("Usage: python aes_cipher.py <file> <key>")
+        return
+
+    fname = args[0]
+    if not os.path.exists(fname) or os.path.isdir(fname) or not os.path.isfile(fname):
+        print("File does not exist or is not a file")
+        return
+
+    key_arg: str = args[1]
+    if len(key_arg) < 16:
+        h = SHA256.new()
+        h.update(key_arg.encode('utf-8'))
+        key = h.digest()[0:16]
+        cipher = AES.new(key, AES.MODE_ECB)
+    else:
+        key = key_arg[0:16]
+        cipher = AES.new(bytes(key.encode('utf-u')), AES.MODE_ECB)
+
+    with open(fname, 'rb') as file:
+        run = True
+        while run:
+            block = file.read(AES.block_size)
+
+            if len(block) != AES.block_size:
+                p = AES.block_size - len(block)
+                block = block + bytes([p]) * p
+                run = False
+
+            cryptogram = cipher.encrypt(block)
+            os.write(1, cryptogram)
+
+
+if __name__ == '__main__':
+    main(sys.argv[1:])
diff --git a/1ano/2semestre/labi/tema01/src/aes_cypher.py b/1ano/2semestre/labi/tema01/src/aes_cypher.py
deleted file mode 100644
index 7946728..0000000
--- a/1ano/2semestre/labi/tema01/src/aes_cypher.py
+++ /dev/null
@@ -1,12 +0,0 @@
-import os
-import sys
-from Crypto.Cipher import AES
-
-
-def main(args=None):
-    cipher = AES.new(os.urandom(16), AES.MODE_ECB)
-    print(cipher.encrypt(b'Hello World!'))
-
-
-if __name__ == '__main__':
-    main(sys.argv[1:])
diff --git a/1ano/2semestre/labi/tema01/src/aes_decipher.py b/1ano/2semestre/labi/tema01/src/aes_decipher.py
new file mode 100644
index 0000000..51f3cea
--- /dev/null
+++ b/1ano/2semestre/labi/tema01/src/aes_decipher.py
@@ -0,0 +1,51 @@
+import os
+import sys
+from Crypto.Cipher import AES
+from Crypto.Hash import SHA256
+
+
+def main(args=None):
+    if len(args) < 2:
+        print("Usage: python aes_decipher.py <file> <key>")
+        return
+
+    fname = args[0]
+    if not os.path.exists(fname) or os.path.isdir(fname) or not os.path.isfile(fname):
+        print("File does not exist or is not a file")
+        return
+
+    key_arg: str = args[1]
+    if len(key_arg) < 16:
+        h = SHA256.new()
+        h.update(key_arg.encode('utf-8'))
+        key = h.digest()[0:16]
+        cipher = AES.new(key, AES.MODE_ECB)
+    else:
+        key = key_arg[0:16]
+        cipher = AES.new(bytes(key.encode('utf-u')), AES.MODE_ECB)
+
+    fsize = os.path.getsize(fname)
+
+    with open(fname, 'rb') as file:
+        run = True
+        while run:
+            block = file.read(AES.block_size)
+            text = cipher.decrypt(block)
+
+            if file.tell() == fsize:
+                has_padding = True
+                i = AES.block_size - text[-1]
+                while has_padding and i < len(text):
+                    has_padding = text[i] == text[-1]
+                    i += 1
+
+                if has_padding:
+                    text = text[:AES.block_size - text[-1]]
+
+                run = False
+
+            os.write(1, text)
+
+
+if __name__ == '__main__':
+    main(sys.argv[1:])