"""A simple helper for "I want to encrypt this chunk of text with a password"
Requires installation of the cryptography library (q2 add_dependency cryptography)
Uses Cryptography's fernet symmetrical encryption algorithm:
https://cryptography.io/en/latest/fernet/
Usage:
.. testcode::
from q2_sdk.core import security
unencrypted_string = "I am in need to encrypting"
secret_key = "supersecret"
encrypted_string = security.encrypt(unencrypted_string, secret_key)
new_unencrypted_string = security.decrypt(encrypted_string, secret_key)
assert unencrypted_string == new_unencrypted_string
"""
import base64
import hashlib
from cryptography.fernet import Fernet
def unpad_str(padded):
pad = padded[-1]
if isinstance(pad, str):
pad = ord(pad)
return padded[:-pad]
def pad_str(data, block_size):
pad = block_size - len(data) % block_size
pad_text = chr(pad) * pad
if isinstance(data, bytes):
pad_text = pad_text.encode()
return data + pad_text
[docs]
def decrypt(encrypted_data: bytes, password: bytes) -> bytes:
"""
Takes in encrypted bytes and a password and returns unencrypted bytes
:param encrypted_data: Encrypted Bytes
:param password: Passphrase to decrypt with
"""
key = encode_password(password)
return Fernet(key).decrypt(encrypted_data)
[docs]
def encrypt(data: bytes | str, password: bytes) -> bytes:
"""
Takes in raw text and a password and returns encrypted bytes
:param data: Raw text
:param password: Passphrase to encrypt with
"""
key = encode_password(password)
if isinstance(data, str):
data = data.encode()
return Fernet(key).encrypt(data)
def encode_password(password):
if isinstance(password, str):
password = password.encode()
digest = hashlib.sha256(password).digest()
encoded_digest = base64.urlsafe_b64encode(digest)
return encoded_digest