From ec69e8ae08aa1fc327ceda592b0dbbc25a68a42a Mon Sep 17 00:00:00 2001 From: hacktobeer Date: Fri, 29 Nov 2024 13:06:55 +0100 Subject: [PATCH] Add create_archive function to create .zip archives. --- openrelik_worker_common/archive_utils.py | 25 ++++++++++++++++++++++++ 1 file changed, 25 insertions(+) diff --git a/openrelik_worker_common/archive_utils.py b/openrelik_worker_common/archive_utils.py index 52eedcf..eabaf60 100644 --- a/openrelik_worker_common/archive_utils.py +++ b/openrelik_worker_common/archive_utils.py @@ -63,3 +63,28 @@ def extract_archive(input_file: dict, output_folder: str, log_file: str) -> str: raise RuntimeError("7zip or tar execution error.") return (command_string, export_folder) + +import shutil +import os + +def create_archive(input_folder_path: str, archive_path: str, delete_input: bool = False): + """ + Zips the contents of a folder into a .zip archive and (optionally) deletes the original folder. + + Args: + input_folder_path (str): Path to the folder to be zipped. + archive_path (str): Path (including filename) where the zip file will be saved. + delete_input (boolean): True if content in folder_path needs to be deleted after archive creation. + + Returns: + str: Path to the created zip file. + """ + if not os.path.isdir(input_folder_path): + raise ValueError(f"The folder {input_folder_path} does not exist.") + + zip_file = shutil.make_archive(archive_path, 'zip', input_folder_path) + + if delete_input: + shutil.rmtree(input_folder_path) + + return zip_file