From a971c2b2360d6b40fa722b4b696a956000603ab9 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Adam=20Hlav=C3=A1=C4=8Dek?= <git@adamhlavacek.com>
Date: Sat, 16 Nov 2024 12:20:44 +0100
Subject: [PATCH] feat(encryption): use dedicated library

---
 encryption.py    | 73 ++++++------------------------------------------
 requirements.txt |  2 +-
 2 files changed, 9 insertions(+), 66 deletions(-)

diff --git a/encryption.py b/encryption.py
index abcb4e5..603bb7a 100644
--- a/encryption.py
+++ b/encryption.py
@@ -1,72 +1,15 @@
-from io import BytesIO
-
-from age.cli import encrypt as age_encrypt, Decryptor as AgeDecryptor, AsciiArmoredInput, AGE_PEM_LABEL
-from age.keys.agekey import AgePrivateKey
+from ssage import SSAGE
+from ssage.backend import SSAGEBackendAge
 
 from config import ENCRYPTION_KEY
 
 
-def __age_key() -> AgePrivateKey:
-    """
-    Get AGE private key
-    :return: AGE private key
-    """
-    return AgePrivateKey.from_private_string(ENCRYPTION_KEY)
-
-
-class __CaptureOnClose(BytesIO):
-    def __init__(self, *args, **kwargs):
-        super().__init__(*args, **kwargs)
-        self.__captured_data = None
-
-    def close(self):
-        self.__captured_data = self.getvalue()
-        super().close()
-
-    @property
-    def captured_data(self):
-        if not self.closed:
-            return self.getvalue()
-
-        data = self.__captured_data
-        self.__captured_data = None
-        return data
-
-
-def encrypt_bytes(data: bytes) -> str:
-    """
-    Encrypt data using AGE encryption
-    :param data: data to encrypt
-    :return: ASCII armored encrypted data
-    """
-    key_public = __age_key().public_key()
-    data_in = BytesIO(data)
-    data_out = __CaptureOnClose()
-
-    age_encrypt(
-        recipients=[key_public.public_string()],
-        infile=data_in,
-        outfile=data_out,
-        ascii_armored=True
-    )
-
-    return data_out.captured_data.decode('ascii')
-
-
-def decrypt_bytes(data: str) -> bytes:
+def get_encryptor() -> SSAGE:
     """
-    Decrypt data using AGE encryption
-    :param data: ASCII armored encrypted data
-    :return: decrypted data
+    Get an encryptor object
+    :return: SSAGE object
     """
-    key = __age_key()
-    data_in = AsciiArmoredInput(AGE_PEM_LABEL, BytesIO(data.encode('ascii')))
-    data_out = __CaptureOnClose()
-
-    with AgeDecryptor([key], data_in) as decryptor:
-        data_out.write(decryptor.read())
-
-    return data_out.captured_data
+    return SSAGE(ENCRYPTION_KEY, authenticate=False, strip=False, backend=SSAGEBackendAge)
 
 
 def encrypt(data: str) -> str:
@@ -75,7 +18,7 @@ def encrypt(data: str) -> str:
     :param data: data to encrypt
     :return: ASCII armored encrypted data
     """
-    return encrypt_bytes(data.encode('utf-8'))
+    return get_encryptor().encrypt(data)
 
 
 def decrypt(data: str) -> str:
@@ -84,4 +27,4 @@ def decrypt(data: str) -> str:
     :param data: ASCII armored encrypted data
     :return: decrypted data
     """
-    return decrypt_bytes(data).decode('utf-8')
+    return get_encryptor().decrypt(data)
diff --git a/requirements.txt b/requirements.txt
index 5afc0fe..171535d 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -23,4 +23,4 @@ smmap2==3.0.1
 SQLAlchemy==2.0.32
 python-dateutil==2.9.0
 PyMySQL==1.1.1
-age==0.5.1
+ssage==1.4.0