grz-cli 0.1.0__py3-none-any.whl

grz_cli/cli.py

@@ -0,0 +1,326 @@
+"""
+CLI module for handling command-line interface operations.
+"""
+
+import logging
+import logging.config
+import sys
+from os import PathLike, sched_getaffinity
+from pathlib import Path
+from tempfile import NamedTemporaryFile
+
+import click
+import platformdirs
+import yaml
+
+from .constants import PACKAGE_ROOT
+from .logging_setup import add_filelogger
+from .models.config import ConfigModel
+from .parser import Worker
+
+log = logging.getLogger(PACKAGE_ROOT + ".cli")
+
+DEFAULT_CONFIG_PATH = Path(platformdirs.user_config_dir("grz-cli")) / "config.yaml"
+
+# Aliases for path types for click options
+# Naming convention: {DIR,FILE}_{Read,Write}_{Exists,Create}
+DIR_R_E = click.Path(
+    exists=True,
+    file_okay=False,
+    dir_okay=True,
+    readable=True,
+    writable=False,
+    resolve_path=True,
+)
+DIR_RW_E = click.Path(
+    exists=True,
+    file_okay=False,
+    dir_okay=True,
+    readable=True,
+    writable=True,
+    resolve_path=True,
+)
+DIR_RW_C = click.Path(
+    exists=False,
+    file_okay=False,
+    dir_okay=True,
+    readable=True,
+    writable=True,
+    resolve_path=True,
+)
+FILE_R_E = click.Path(
+    exists=True, file_okay=True, dir_okay=False, readable=True, resolve_path=True
+)
+
+submission_dir = click.option(
+    "--submission-dir",
+    metavar="PATH",
+    type=DIR_R_E,
+    required=True,
+    help="Path to the submission directory containing 'metadata/', 'files/', 'encrypted_files/' and 'logs/' directories",
+)
+
+config_file = click.option(
+    "--config-file",
+    metavar="STRING",
+    type=FILE_R_E,
+    required=False,
+    default=DEFAULT_CONFIG_PATH,
+    help="Path to config file",
+)
+
+threads = click.option(
+    "--threads",
+    default=min(len(sched_getaffinity(0)), 4),
+    type=int,
+    show_default=True,
+    help="Number of threads to use for parallel operations",
+)
+
+submission_id = click.option(
+    "--submission-id",
+    required=True,
+    type=str,
+    metavar="STRING",
+    help="S3 submission prefix (corresponds to the tanG of a submission)",
+)
+
+output_dir = click.option(
+    "--output-dir",
+    metavar="PATH",
+    type=DIR_RW_E,
+    required=True,
+    default=None,
+    help="Path to the target submission output directory",
+)
+
+
+class OrderedGroup(click.Group):
+    """
+    A click Group that keeps track of the order in which commands are added.
+    """
+
+    def list_commands(self, ctx):
+        """Return the list of commands in the order they were added."""
+        return list(self.commands.keys())
+
+
+@click.group(
+    cls=OrderedGroup,
+    help="Validate, encrypt, decrypt and upload submissions to a GRZ/GDC.",
+)
+@click.version_option(version="0.1", prog_name="grz-cli")
+@click.option("--log-file", metavar="FILE", type=str, help="Path to log file")
+@click.option(
+    "--log-level",
+    default="INFO",
+    type=click.Choice(["DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"]),
+    help="Set the log level (default: INFO)",
+)
+def cli(log_file: str | None = None, log_level: str = "INFO"):
+    """
+    Command-line interface function for setting up logging.
+
+    :param log_file: Path to the log file. If provided, a file logger will be added.
+    :param log_level: Log level for the logger. It should be one of the following:
+                       DEBUG, INFO, WARNING, ERROR, CRITICAL.
+    """
+    if log_file:
+        add_filelogger(
+            log_file,
+            log_level.upper(),
+        )  # Add file logger
+
+    # show only time and log level in STDOUT
+    logging.basicConfig(
+        format="%(asctime)s - %(levelname)s - %(message)s",
+    )
+
+    # set the log level for this package
+    logging.getLogger(PACKAGE_ROOT).setLevel(log_level.upper())
+
+    log.debug("Logging setup complete.")
+
+
+@cli.command()
+@submission_dir
+def validate(submission_dir):
+    """
+    Validate the submission.
+
+    This validates the submission by checking its checksums, as well as performing basic sanity checks on the supplied metadata.
+    Must be executed before calling `encrypt` and `upload`.
+    """
+    log.info("Starting validation...")
+
+    submission_dir = Path(submission_dir)
+
+    worker_inst = Worker(
+        metadata_dir=submission_dir / "metadata",
+        files_dir=submission_dir / "files",
+        log_dir=submission_dir / "logs",
+        encrypted_files_dir=submission_dir / "encrypted_files",
+    )
+    worker_inst.validate()
+
+    log.info("Validation finished!")
+
+
+@cli.command()
+@submission_dir
+@config_file
+def encrypt(
+    submission_dir,
+    config_file,
+):
+    """
+    Encrypt a submission.
+
+    Encryption is done with the recipient's public key.
+    Sub-folders 'encrypted_files' and 'logs' are created within the submission directory.
+    """
+    config = read_config(config_file)
+
+    submitter_privkey_path = config.submitter_private_key_path
+    if submitter_privkey_path == "":
+        submitter_privkey_path = None
+
+    log.info("Starting encryption...")
+
+    submission_dir = Path(submission_dir)
+
+    worker_inst = Worker(
+        metadata_dir=submission_dir / "metadata",
+        files_dir=submission_dir / "files",
+        log_dir=submission_dir / "logs",
+        encrypted_files_dir=submission_dir / "encrypted_files",
+    )
+    if config.grz_public_key:
+        with NamedTemporaryFile("w") as f:
+            f.write(config.grz_public_key)
+            f.flush()
+            worker_inst.encrypt(
+                f.name,
+                submitter_private_key_path=submitter_privkey_path,
+            )
+    else:
+        worker_inst.encrypt(
+            config.grz_public_key_path,
+            submitter_private_key_path=submitter_privkey_path,
+        )
+
+    log.info("Encryption successful!")
+
+
+@cli.command()
+@submission_dir
+@config_file
+def decrypt(
+    submission_dir,
+    config_file,
+):
+    """
+    Decrypt a submission.
+
+    Decrypting a submission requires the _private_ key of the original recipient.
+    """
+    config = read_config(config_file)
+
+    grz_privkey_path = config.grz_private_key_path
+    if not grz_privkey_path:
+        log.error("GRZ private key path is required for decryption.")
+        sys.exit(1)
+
+    log.info("Starting encryption...")
+
+    submission_dir = Path(submission_dir)
+
+    worker_inst = Worker(
+        metadata_dir=submission_dir / "metadata",
+        files_dir=submission_dir / "files",
+        log_dir=submission_dir / "logs",
+        encrypted_files_dir=submission_dir / "encrypted_files",
+    )
+    worker_inst.decrypt(grz_privkey_path)
+
+    log.info("Encryption successful!")
+
+
+@cli.command()
+@submission_dir
+@config_file
+@threads
+def upload(
+    submission_dir,
+    config_file,
+    threads,
+):
+    """
+    Upload a submission to a GRZ/GDC.
+    """
+    config = read_config(config_file)
+
+    log.info("Starting upload...")
+
+    submission_dir = Path(submission_dir)
+
+    worker_inst = Worker(
+        metadata_dir=submission_dir / "metadata",
+        files_dir=submission_dir / "files",
+        log_dir=submission_dir / "logs",
+        encrypted_files_dir=submission_dir / "encrypted_files",
+        threads=threads,
+    )
+    worker_inst.upload(config)
+
+    log.info("Upload finished!")
+
+
+@cli.command()
+@submission_id
+@output_dir
+@config_file
+@threads
+def download(
+    submission_id,
+    output_dir,
+    config_file,
+    threads,
+):
+    """
+    Download a submission from a GRZ.
+
+    Downloaded metadata is stored within the `metadata` sub-folder of the submission output directory.
+    Downloaded files are stored within the `encrypted_files` sub-folder of the submission output directory.
+    """
+    config = read_config(config_file)
+
+    log.info("Starting download...")
+
+    submission_dir_path = Path(output_dir)
+    if not submission_dir_path.is_dir():
+        log.debug("Creating submission directory %s", submission_dir_path)
+        submission_dir_path.mkdir(mode=0o770, parents=False, exist_ok=False)
+
+    worker_inst = Worker(
+        metadata_dir=submission_dir_path / "metadata",
+        files_dir=submission_dir_path / "files",
+        log_dir=submission_dir_path / "logs",
+        encrypted_files_dir=submission_dir_path / "encrypted_files",
+        threads=threads,
+    )
+    worker_inst.download(config, submission_id)
+
+    log.info("Download finished!")
+
+
+def read_config(config_path: str | PathLike) -> ConfigModel:
+    """Reads the configuration file and validates it against the schema."""
+    with open(config_path, encoding="utf-8") as f:
+        config = ConfigModel(**yaml.safe_load(f))
+
+    return config
+
+
+if __name__ == "__main__":
+    cli()

grz_cli/constants.py

@@ -0,0 +1,37 @@
+"""Constants for logging configuration, JSON schema validation, and other settings."""
+
+PACKAGE_ROOT = "grz_cli"
+
+LOGGING_FORMAT = "%(asctime)s [%(levelname)s] %(name)s: %(message)s"
+LOGGING_DATEFMT = "%Y-%m-%d %I:%M %p"
+LOGGING_CONFIG = {
+    "version": 1,
+    "disable_existing_loggers": True,
+    "formatters": {
+        "standard": {
+            "format": LOGGING_FORMAT,
+            "datefmt": LOGGING_DATEFMT,
+        },
+    },
+    "handlers": {
+        "default": {
+            "level": "INFO",
+            "formatter": "standard",
+            "class": "logging.StreamHandler",
+            "stream": "ext://sys.stdout",  # Default is stderr
+        },
+    },
+    "loggers": {
+        "": {  # root logger
+            "handlers": ["default"],
+            "level": "WARNING",
+            "propagate": False,
+        },
+        PACKAGE_ROOT: {"handlers": ["default"], "level": "INFO", "propagate": False},
+        "__main__": {  # if __name__ == '__main__'
+            "handlers": ["default"],
+            "level": "DEBUG",
+            "propagate": False,
+        },
+    },
+}

grz_cli/download.py

@@ -0,0 +1,236 @@
+"""Module for uploading encrypted submissions to a remote storage"""
+
+from __future__ import annotations
+
+import logging
+import math
+import re
+from os import PathLike
+from pathlib import Path
+from typing import TYPE_CHECKING
+
+import boto3  # type: ignore[import-untyped]
+import botocore.handlers  # type: ignore[import-untyped]
+from boto3 import client as boto3_client  # type: ignore[import-untyped]
+from boto3.s3.transfer import S3Transfer, TransferConfig  # type: ignore[import-untyped]
+from botocore.config import Config as Boto3Config  # type: ignore[import-untyped]
+from tqdm.auto import tqdm
+
+from .models.config import ConfigModel
+from .progress_logging import FileProgressLogger
+
+MULTIPART_THRESHOLD = 8 * 1024 * 1024  # 8MiB, boto3 default
+MULTIPART_CHUNKSIZE = 8 * 1024 * 1024  # 8MiB, boto3 default
+MULTIPART_MAX_CHUNKS = 1000  # CEPH S3 limit, AWS limit is 10000
+
+if TYPE_CHECKING:
+    from .parser import EncryptedSubmission
+
+log = logging.getLogger(__name__)
+
+# see discussion: https://github.com/boto/boto3/discussions/4251 to accept bucket names with ":" in the name
+botocore.handlers.VALID_BUCKET = re.compile(r"^[:a-zA-Z0-9.\-_]{1,255}$")  # type: ignore[import-untyped]
+
+
+class DownloadError(Exception):
+    """Exception raised when an upload fails"""
+
+    pass
+
+
+class S3BotoDownloadWorker:
+    """Implementation of an upload worker using boto3 for S3"""
+
+    __log = log.getChild("S3BotoDownloadWorker")
+
+    def __init__(
+        self,
+        config: ConfigModel,
+        status_file_path: str | PathLike,
+        threads: int = 1,
+    ):
+        """
+        An download manager for S3 storage
+
+        :param config: The configuration model
+        :param status_file_path: The path to the status file
+        """
+        super().__init__()
+
+        self._status_file_path = Path(status_file_path)
+        self._config = config
+        self._threads = threads
+
+        self._init_s3_client()
+
+    def _init_s3_client(self):
+        # if user specifies empty strings, this might be an issue
+        def empty_str_to_none(string: str | None) -> str | None:
+            if string == "" or string is None:
+                return None
+            else:
+                return string
+
+        # configure proxies if proxy_url is defined
+        proxy_url = empty_str_to_none(self._config.s3_options.proxy_url)
+        if proxy_url is not None:
+            config = Boto3Config(proxies={"http": proxy_url, "https": proxy_url})
+        else:
+            config = None
+
+        # Initialize S3 client for uploading
+        self._s3_client: boto3.session.Session.client = boto3_client(
+            service_name="s3",
+            region_name=empty_str_to_none(self._config.s3_options.region_name),
+            api_version=empty_str_to_none(self._config.s3_options.api_version),
+            use_ssl=self._config.s3_options.use_ssl,
+            endpoint_url=empty_str_to_none(str(self._config.s3_options.endpoint_url)),
+            aws_access_key_id=empty_str_to_none(self._config.s3_options.access_key),
+            aws_secret_access_key=empty_str_to_none(self._config.s3_options.secret),
+            aws_session_token=empty_str_to_none(self._config.s3_options.session_token),
+            config=config,
+        )
+
+    def prepare_download(
+        self,
+        metadata_dir: Path,
+        encrypted_files_dir: Path,
+        log_dir: Path,
+    ):
+        """
+        Prepare the download of an encrypted submission
+
+        :param metadata_dir: Path to the metadata directory
+        :param encrypted_files_dir: Path to the encrypted_files directory
+        :param log_dir: Path to the logs directory
+        """
+        for dir_path in [metadata_dir, encrypted_files_dir, log_dir]:
+            if not dir_path.exists():
+                self.__log.debug("Creating directory: %s", dir_path)
+                dir_path.mkdir(parents=False, exist_ok=False)  # Create the directories
+            else:
+                self.__log.debug("Directory exists: %s", dir_path)
+
+    def download_metadata(
+        self,
+        submission_id: str,
+        metadata_dir: Path,
+        metadata_file_name: str = "metadata.json",
+    ):
+        """
+        Download the metadata.json
+
+        :param submission_id: submission folder on S3 structure
+        :param metadata_dir: Path of the metadir folder
+        :param metadata_file_name: name of the metadata.json
+        """
+        metadata_key = str(Path(submission_id) / metadata_dir.name / metadata_file_name)
+        metadata_file_path = metadata_dir / metadata_file_name
+
+        try:
+            self.download_file(str(metadata_file_path), metadata_key, 10000)
+        except Exception as e:
+            self.__log.error("Download failed for '%s'", str(metadata_key))
+
+            raise e
+
+    def download_file(self, local_file_path, s3_object_id, filesize):
+        """
+        Upload a single file to the specified object ID
+        :param local_file_path: Path to the file to upload
+        :param s3_object_id: Remote S3 object ID under which the file should be stored
+        :param filesize: size of the file
+        """
+        # self.__log.info(f"Download {s3_object_id} to {local_file_path}...")
+
+        chunksize = (
+            math.ceil(filesize / MULTIPART_MAX_CHUNKS)
+            if filesize / MULTIPART_CHUNKSIZE > MULTIPART_MAX_CHUNKS
+            else MULTIPART_CHUNKSIZE
+        )
+        self.__log.debug(
+            f"Using a chunksize of: {chunksize / 1024**2}MiB, results in {filesize/chunksize} chunks"
+        )
+
+        config = TransferConfig(
+            multipart_threshold=MULTIPART_THRESHOLD,
+            multipart_chunksize=chunksize,
+            max_concurrency=self._threads,
+        )
+
+        transfer = S3Transfer(self._s3_client, config)
+        progress_bar = tqdm(
+            total=filesize, unit="B", unit_scale=True, unit_divisor=1024
+        )
+        transfer.download_file(
+            self._config.s3_options.bucket,
+            s3_object_id,
+            local_file_path,
+            callback=lambda bytes_transferred: progress_bar.update(bytes_transferred),
+        )
+
+    def download(self, submission_id: str, encrypted_submission: EncryptedSubmission):
+        """
+        Upload an encrypted submission
+        :param encrypted_submission: The encrypted submission to upload
+        """
+        progress_logger = FileProgressLogger(self._status_file_path)
+
+        encrypted_files_key_prefix = f"{submission_id}/files/"
+
+        response = self._s3_client.list_objects_v2(
+            Bucket=self._config.s3_options.bucket, Prefix=encrypted_files_key_prefix
+        )
+
+        if "Contents" in response:
+            for file in response["Contents"]:
+                if file["Size"] == 0:
+                    continue
+                else:
+                    file_key = file["Key"]
+                    file_path = Path(file_key).relative_to(encrypted_files_key_prefix)
+                    full_path = encrypted_submission.encrypted_files_dir / file_path
+                    if full_path not in encrypted_submission.encrypted_files:
+                        raise DownloadError(
+                            f"File {file_path} not listed in metadata.json"
+                        )
+                    file_metadata = encrypted_submission.encrypted_files[full_path]
+                    logged_state = progress_logger.get_state(full_path, file_metadata)
+
+                    if (logged_state is None) or not logged_state.get(
+                        "download_successful", False
+                    ):
+                        self.__log.info(
+                            "Download file: '%s' -> '%s'",
+                            file_key,
+                            str(full_path),
+                        )
+
+                        try:
+                            full_path.parent.mkdir(
+                                mode=0o770, parents=True, exist_ok=True
+                            )
+                            self.download_file(str(full_path), file_key, file["Size"])
+
+                            self.__log.info(f"Download complete for {str(full_path)}. ")
+                            progress_logger.set_state(
+                                full_path,
+                                file_metadata,
+                                state={"download_successful": True},
+                            )
+                        except Exception as e:
+                            self.__log.error("Download failed for '%s'", str(full_path))
+
+                            progress_logger.set_state(
+                                full_path,
+                                file_metadata,
+                                state={"download_successful": False, "error": str(e)},
+                            )
+
+                            raise e
+                    else:
+                        self.__log.info(
+                            "File '%s' already downloaded (at '%s')",
+                            file_key,
+                            str(full_path),
+                        )