From e3f90dacc656a97ced0860971d1c68254a48d69c Mon Sep 17 00:00:00 2001 From: gonzcm Date: Sat, 4 May 2024 18:47:51 +0200 Subject: [PATCH] feat: Add encrypted backup dump and restore functionlity #265 --- src/core/utils/security.py | 57 +++++++++++++++++ src/modules/shared/backup/__init__.py | 0 src/modules/shared/backup/controller.py | 24 +++++++ src/modules/shared/backup/router.py | 31 +++++++++ src/modules/shared/backup/service.py | 84 +++++++++++++++++++++++++ 5 files changed, 196 insertions(+) create mode 100644 src/modules/shared/backup/__init__.py create mode 100644 src/modules/shared/backup/controller.py create mode 100644 src/modules/shared/backup/router.py create mode 100644 src/modules/shared/backup/service.py diff --git a/src/core/utils/security.py b/src/core/utils/security.py index b3c0eba..e101b2e 100644 --- a/src/core/utils/security.py +++ b/src/core/utils/security.py @@ -1,9 +1,15 @@ from datetime import datetime, timedelta +import json +import os from typing import Union, Any import bcrypt from jose import jwt +from cryptography.hazmat.primitives import padding +from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes +from cryptography.hazmat.backends import default_backend from src.core.config import settings +from src.core.database.backup import BackupEncoder def create_access_token( @@ -56,3 +62,54 @@ def verify_password(password: str, hashed_password: str) -> bool: password_byte_enc = password.encode('utf-8') hashed_password_bytes = hashed_password.encode('utf-8') return bcrypt.checkpw(password_byte_enc, hashed_password_bytes) + + +def generate_salt() -> bytes: + return os.urandom(16) + + +def derive_key(password: str, salt: bytes) -> bytes: + key = bcrypt.kdf(password.encode(), salt, + desired_key_bytes=32, rounds=100) + return key + + +def pad_data(data): + padder = padding.PKCS7(128).padder() + padded_data = padder.update(data) + padder.finalize() + return padded_data + + +def unpad_data(data): + unpadder = padding.PKCS7(128).unpadder() + unpadded_data = unpadder.update(data) + unpadder.finalize() + return unpadded_data + + +def encrypt_data(data, key: bytes): + iv = os.urandom(16) + + cipher = Cipher(algorithms.AES(key), modes.CBC(iv), + backend=default_backend()) + + encryptor = cipher.encryptor() + + padded_data = pad_data(data) + + encrypted_data = encryptor.update(padded_data) + encryptor.finalize() + + return iv, encrypted_data + + +def decrypt_data(encrypted_data, key, iv): + cipher = Cipher(algorithms.AES(key), modes.CBC(iv), + backend=default_backend()) + + decryptor = cipher.decryptor() + + decrypted_data = decryptor.update(encrypted_data) + decryptor.finalize() + + # Remove padding after decryption + unpadded_data = unpad_data(decrypted_data) + + return unpadded_data diff --git a/src/modules/shared/backup/__init__.py b/src/modules/shared/backup/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/modules/shared/backup/controller.py b/src/modules/shared/backup/controller.py new file mode 100644 index 0000000..309f80b --- /dev/null +++ b/src/modules/shared/backup/controller.py @@ -0,0 +1,24 @@ +from fastapi import status, HTTPException + +from src.core.deps import DataBaseDep +from src.modules.shared.backup import service + + +async def generate_backup_controller(db: DataBaseDep, password: str): + try: + return await service.generate_backup_service(db, password) + except Exception as e: + raise HTTPException( + status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, + detail=f"Internal Server Error {e}" + ) from e + + +async def restore_backup_controller(db: DataBaseDep, password: str, file): + try: + return await service.restore_backup_service(db, password, file) + except Exception as e: + raise HTTPException( + status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, + detail=f"Internal Server Error {e}" + ) from e diff --git a/src/modules/shared/backup/router.py b/src/modules/shared/backup/router.py new file mode 100644 index 0000000..17866d6 --- /dev/null +++ b/src/modules/shared/backup/router.py @@ -0,0 +1,31 @@ +from fastapi import APIRouter, File, UploadFile, status + +from src.core.deps import DataBaseDep +from src.modules.shared.backup import controller + + +router = APIRouter(tags=['Backup']) + + +@router.get('/dump', status_code=status.HTTP_200_OK, responses={ + status.HTTP_200_OK: {'description': 'Successful Response'}, + status.HTTP_500_INTERNAL_SERVER_ERROR: { + 'description': 'Internal Server Error'} +}) +async def generate_backup(db: DataBaseDep, password: str): + """ + Generates a backup of the database. + """ + return await controller.generate_backup_controller(db, password) + + +@router.post('/restore', status_code=status.HTTP_200_OK, responses={ + status.HTTP_200_OK: {'description': 'Successful Response'}, + status.HTTP_500_INTERNAL_SERVER_ERROR: { + 'description': 'Internal Server Error'} +}) +async def restore_backup(db: DataBaseDep, password: str, file: UploadFile = File(...)): + """ + Restores a backup of the database. + """ + return await controller.restore_backup_controller(db, password, file) diff --git a/src/modules/shared/backup/service.py b/src/modules/shared/backup/service.py new file mode 100644 index 0000000..428c06d --- /dev/null +++ b/src/modules/shared/backup/service.py @@ -0,0 +1,84 @@ +import json +import os +import tarfile +import tempfile + +from fastapi import HTTPException, Response +from fastapi.responses import JSONResponse +from src.core.database.backup import BackupEncoder, dump_to_json, populate_from_json +from src.core.deps import DataBaseDep +from src.core.utils.security import decrypt_data, derive_key, encrypt_data, generate_salt + + +async def generate_backup_service(db: DataBaseDep, password: str): + data = await dump_to_json(db) + + print("Data dumped") + + with tempfile.TemporaryDirectory() as tmp_dir: + json_file_path = os.path.join(tmp_dir, "backup.json") + with open(json_file_path, "w") as json_file: + json.dump(data, json_file, indent=4, cls=BackupEncoder) + + print("Data dumped to JSON file") + + compressed_file_path = os.path.join(tmp_dir, "backup.tar.gz") + with tarfile.open(compressed_file_path, "w:gz") as tar: + tar.add(json_file_path, arcname="backup.json") + + print("Data compressed") + + with open(compressed_file_path, "rb") as compressed_file: + compressed_data = compressed_file.read() + + salt = generate_salt() + key = derive_key(password, salt) + + print("Key derived") + + iv, encrypted_data = encrypt_data(compressed_data, key) + + print("Data encrypted") + + combined_data = salt + iv + encrypted_data + + return Response( + content=combined_data, + media_type="application/octet-stream", + headers={"Content-Disposition": "attachment; filename=backup.enc"} + ) + + +async def restore_backup_service(db: DataBaseDep, password: str, file): + with tempfile.TemporaryDirectory() as tmp_dir: + file_path = os.path.join(tmp_dir, file.filename) + with open(file_path, "wb") as tmp_file: + tmp_file.write(await file.read()) + + with open(file_path, "rb") as encrypted_file: + combined_data = encrypted_file.read() + + salt = combined_data[:16] + combined_data = combined_data[16:] + + iv = combined_data[:16] + encrypted_data = combined_data[16:] + + key = derive_key(password, salt) + + decrypted_data = decrypt_data(encrypted_data, key, iv) + + decrypted_file_path = os.path.join( + tmp_dir, "decrypted_data.tar.gz") + with open(decrypted_file_path, "wb") as decrypted_file: + decrypted_file.write(decrypted_data) + + with tarfile.open(decrypted_file_path, mode="r:gz") as tar: + tar.extractall(tmp_dir) + + with open(os.path.join(tmp_dir, "backup.json"), "r") as json_file: + restored_data = json.load(json_file) + + await populate_from_json(db, restored_data) + + return {"message": "Data restored successfully"}