Skip to content

Commit

Permalink
feat(encryption): use dedicated library
Browse files Browse the repository at this point in the history
  • Loading branch information
esoadamo committed Nov 16, 2024
1 parent 522d45b commit a971c2b
Show file tree
Hide file tree
Showing 2 changed files with 9 additions and 66 deletions.
73 changes: 8 additions & 65 deletions encryption.py
Original file line number Diff line number Diff line change
@@ -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:
Expand All @@ -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:
Expand All @@ -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)
2 changes: 1 addition & 1 deletion requirements.txt
Original file line number Diff line number Diff line change
Expand Up @@ -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

0 comments on commit a971c2b

Please sign in to comment.