From 1bb8aff9da30c602bb756a76d0e7e021e78637d1 Mon Sep 17 00:00:00 2001 From: Juan Carlos Jose Camacho Date: Tue, 19 Sep 2023 11:23:23 -0600 Subject: [PATCH] DH-4603 Decrypt file content if it is encrypted --- dataherald/utils/s3.py | 18 +++++++++++++++--- 1 file changed, 15 insertions(+), 3 deletions(-) diff --git a/dataherald/utils/s3.py b/dataherald/utils/s3.py index 3dc6c25c..e185e482 100644 --- a/dataherald/utils/s3.py +++ b/dataherald/utils/s3.py @@ -1,6 +1,8 @@ import boto3 +from cryptography.fernet import InvalidToken from dataherald.config import Settings +from dataherald.utils.encrypt import FernetEncrypt class S3: @@ -8,14 +10,24 @@ def __init__(self): self.settings = Settings() def download(self, path: str) -> str: + fernet_encrypt = FernetEncrypt() path = path.split("/") s3_client = boto3.client( "s3", aws_access_key_id=self.settings.s3_aws_access_key_id, aws_secret_access_key=self.settings.s3_aws_secret_access_key, ) - + file_location = f"tmp/{path[-1]}" s3_client.download_file( - Bucket=path[2], Key=f"{path[-1]}", Filename=f"tmp/{path[-1]}" + Bucket=path[2], Key=f"{path[-1]}", Filename=file_location ) - return f"tmp/{path[-1]}" + # Decrypt file content if it is encrypted + try: + with open(file_location) as file_object: + decrypted_content = fernet_encrypt.decrypt(file_object.read()) + with open(file_location, "w") as file_object: + file_object.write(decrypted_content) + except InvalidToken: + pass + + return file_location