Added new Encrypted column type for AES256 encrypted columns

Author: maximebf

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "sqlorm-py"
-version = "0.2.1"
+version = "0.2.2"
 description = "A new kind or ORM that do not abstract away your database or SQL queries."
 authors = ["Maxime Bouroumeau-Fuseau <[email protected]>"]
 readme = "README.md"
@@ -13,6 +13,7 @@ python = "^3.10"
 blinker = "^1.8.2"
 psycopg = { extras = ["binary"], version = "^3.1.18", optional = true }
 mysql-connector-python = { version = "^8.3.0", optional = true }
+pycryptodome = { version = "^3.20.0", optional = true }
 
 [tool.poetry.group.dev.dependencies]
 pytest = "^8.0.0"

sqlorm/types.py renamed to sqlorm/types/__init__.py

@@ -57,6 +57,8 @@ def __str__(self):
     }
 )
 
+from .encrypted import Encrypted
+
 __all__ = (
     "SQLType",
     "Integer",
@@ -68,4 +70,5 @@ def __str__(self):
     "DateTime",
     "JSON",
     "Pickle",
+    "Encrypted"
 )

sqlorm/types/encrypted.py

@@ -0,0 +1,33 @@
+import base64
+from Crypto.Cipher import AES
+from Crypto import Random
+import os
+from . import SQLType
+
+
+class Encrypted(SQLType):
+    block_size = 16
+
+    def __init__(self, key):
+        super().__init__("text", self.decrypt, self.encrypt)
+        self.key = key
+
+    def encrypt(self, raw):
+        raw = self.pad(raw).encode()
+        iv = Random.new().read(AES.block_size)
+        key = self.key() if callable(self.key) else self.key
+        cipher = AES.new(key, AES.MODE_CBC, iv)
+        return base64.b64encode(iv + cipher.encrypt(raw))
+
+    def decrypt(self, enc):
+        enc = base64.b64decode(enc)
+        iv = enc[:16]
+        key = self.key() if callable(self.key) else self.key
+        cipher = AES.new(key, AES.MODE_CBC, iv)
+        return self.unpad(cipher.decrypt(enc[16:]).decode())
+
+    def pad(self, s):
+        return s + (self.block_size - len(s) % self.block_size) * chr(self.block_size - len(s) % self.block_size)
+    
+    def unpad(self, s):
+        return s[:-ord(s[len(s)-1:])]