81 lines
2.0 KiB
Python
81 lines
2.0 KiB
Python
# NOTE: Cryptography import
|
|
from Crypto.Cipher import AES
|
|
|
|
# NOTE: Generic imports
|
|
import traceback
|
|
import logging
|
|
|
|
|
|
logger = logging.getLogger("__main__." + __name__)
|
|
|
|
|
|
class SteelPetal:
|
|
"""
|
|
Cryptography utility for encrypting files
|
|
"""
|
|
|
|
def __init__(self, key: str, nonce=None, testData=None):
|
|
"""
|
|
Parameters
|
|
----------
|
|
key: str
|
|
User's plaintext key
|
|
|
|
nonce
|
|
Cryptographic artifact we can use to reinitialize cryptographic operations
|
|
"""
|
|
try:
|
|
if nonce is None:
|
|
self.cipher = AES.new(self.pad(key).encode("utf-8"), AES.MODE_GCM)
|
|
self.nonce = self.cipher.nonce
|
|
else:
|
|
self.cipher = AES.new(
|
|
self.pad(key).encode("utf-8"), AES.MODE_GCM, nonce=nonce
|
|
)
|
|
self.nonce = nonce
|
|
if testData is not None:
|
|
try:
|
|
self.cipher.decrypt(testData)
|
|
except Exception:
|
|
logger.error(traceback.format_exc())
|
|
return False
|
|
except Exception:
|
|
logger.error(traceback.format_exc())
|
|
|
|
def pad(self, key: str):
|
|
"""
|
|
Pad key to make it usable
|
|
|
|
key: str
|
|
User's plain text key
|
|
"""
|
|
BS = AES.block_size
|
|
key = key + (BS - len(key) % BS) * chr(BS - len(key) % BS)
|
|
return key
|
|
|
|
def encrypt(self, data: bytes):
|
|
"""
|
|
Encrypt binary data
|
|
|
|
data: bytes
|
|
Data to encrypt
|
|
"""
|
|
try:
|
|
return self.cipher.encrypt_and_digest(data)
|
|
except Exception:
|
|
logger.error(traceback.format_exec())
|
|
return False
|
|
|
|
def decrypt(self, data: bytes):
|
|
"""
|
|
Decrypt encrypted binary data
|
|
|
|
data: bytes
|
|
Data to decrypt
|
|
"""
|
|
try:
|
|
return self.cipher.decrypt(data)
|
|
except Exception:
|
|
logger.error(traceback.format_exec())
|
|
return False
|