path: root/src/sec_certs/cli.py

#!/usr/bin/env python3
from __future__ import annotations

import logging
import sys
from collections.abc import Callable
from dataclasses import dataclass, field
from datetime import datetime
from pathlib import Path

import click
from pydantic import ValidationError

from sec_certs.configuration import config
from sec_certs.dataset.cc import CCDataset
from sec_certs.dataset.dataset import Dataset
from sec_certs.dataset.fips import FIPSDataset
from sec_certs.dataset.protection_profile import ProtectionProfileDataset
from sec_certs.utils.helpers import warn_if_missing_poppler, warn_if_missing_tesseract

logger = logging.getLogger(__name__)

EXIT_CODE_NOK: int = 1
EXIT_CODE_OK: int = 0

DEFAULT_OUTPUTPATH: Path = Path("./dataset").resolve()


@dataclass
class ProcessingStep:
    name: str
    processing_function_name: str
    preconditions: list[str] = field(default_factory=list)
    precondition_error_msg: str | None = field(default=None)
    pre_callback_func: Callable | None = field(default=None)

    def __post_init__(self) -> None:
        for condition in self.preconditions:
            if not hasattr(Dataset.DatasetInternalState, condition):
                raise ValueError(f"Precondition attribute {condition} is not member of `Dataset.DatasetInternalState`.")

    def run(self, dset: Dataset) -> None:
        for condition in self.preconditions:
            if not getattr(dset.state, condition):
                err_msg = (
                    self.precondition_error_msg
                    if self.precondition_error_msg
                    else f"Error, precondition to run {self.name} not met, exiting."
                )
                click.echo(err_msg, err=True)
                sys.exit(EXIT_CODE_NOK)
        if self.pre_callback_func:
            self.pre_callback_func()

        getattr(dset, self.processing_function_name)()


def warn_missing_libs():
    warn_if_missing_poppler()
    warn_if_missing_tesseract()


FRAMEWORK_TO_CONSTRUCTOR: dict[str, type[Dataset]] = {
    "cc": CCDataset,
    "fips": FIPSDataset,
    "pp": ProtectionProfileDataset,
}


def build_or_load_dataset(
    framework: str,
    inputpath: Path | None,
    to_build: bool,
    outputpath: Path | None,
) -> Dataset:
    constructor: type[Dataset] = FRAMEWORK_TO_CONSTRUCTOR[framework]
    dset: Dataset

    if to_build:
        if not outputpath:
            outputpath = DEFAULT_OUTPUTPATH
        if inputpath:
            print(
                f"Warning: you wanted to build a dataset but you provided one in JSON -- that will be ignored. New one will be constructed at: {outputpath}"
            )
        dset = constructor(
            certs={},
            root_dir=outputpath,
            name=framework + "_dataset",
            description=f"Full {framework} dataset snapshot {datetime.now().date()}",
        )
        dset.get_certs_from_web()
    else:
        if not inputpath:
            click.echo(
                "Error: If you do not use 'build' action, you must provide --input parameter to point to an existing dataset.",
                err=True,
            )
            sys.exit(EXIT_CODE_NOK)

        dset = constructor.from_json(inputpath)

        if outputpath and dset.root_dir != outputpath:
            print(
                "Warning: you provided both input and output paths. The dataset from input path will get copied to output path."
            )
            dset.copy_dataset(outputpath)

    return dset


steps = [
    ProcessingStep(
        "process-aux-dsets",
        "process_auxiliary_datasets",
        preconditions=["meta_sources_parsed"],
        precondition_error_msg="Error: You want to process the auxiliary datasets, but the data from cert. framework website was not parsed. You must use 'build' action first.",
        pre_callback_func=None,
    ),
    ProcessingStep(
        "download",
        "download_all_artifacts",
        preconditions=["meta_sources_parsed"],
        precondition_error_msg="Error: You want to download all artifacts, but the data from the cert. framework website was not parsed. You must use 'build' action first.",
        pre_callback_func=None,
    ),
    ProcessingStep(
        "convert",
        "convert_all_pdfs",
        preconditions=["artifacts_downloaded"],
        precondition_error_msg="Error: You want to convert pdfs -> txt, but the pdfs were not downloaded. You must use 'download' action first.",
        pre_callback_func=warn_missing_libs,
    ),
    ProcessingStep(
        "analyze",
        "analyze_certificates",
        preconditions=["pdfs_converted", "auxiliary_datasets_processed"],
        precondition_error_msg="Error: You want to process txt documents of certificates, but pdfs were not converted. You must use 'convert' action first.",
        pre_callback_func=None,
    ),
]


@click.command()
@click.argument(
    "framework",
    required=True,
    nargs=1,
    type=click.Choice(["cc", "fips", "pp"], case_sensitive=False),
)
@click.argument(
    "actions",
    required=True,
    nargs=-1,
    type=click.Choice(["all", "build", "process-aux-dsets", "download", "convert", "analyze"], case_sensitive=False),
)
@click.option(
    "-o",
    "--output",
    "outputpath",
    type=click.Path(file_okay=False, dir_okay=True, writable=True, readable=True, resolve_path=True),
    help="Path where the output of the experiment will be stored. May overwrite existing content.",
)
@click.option(
    "-c",
    "--config",
    "configpath",
    default=None,
    type=click.Path(file_okay=True, dir_okay=False, writable=True, readable=True),
    help="Path to your own config yaml file that will override the default config.",
)
@click.option(
    "-i",
    "--input",
    "inputpath",
    type=click.Path(file_okay=True, dir_okay=False, writable=True, readable=True),
    help="If set, the actions will be performed on a dataset loaded from JSON from the input path.",
)
@click.option("-q", "--quiet", is_flag=True, help="If set, will not print to stdout")
def main(
    framework: str,
    actions: list[str],
    outputpath: Path | None,
    configpath: Path | None,
    inputpath: Path | None,
    quiet: bool,
):
    try:
        if configpath:
            try:
                config.load_from_yaml(configpath)
            except FileNotFoundError:
                click.echo("Error: Bad path to configuration file", err=True)
                sys.exit(EXIT_CODE_NOK)
            except (ValueError, ValidationError) as e:
                click.echo(f"Error: Bad format of configuration file: {e}", err=True)
                sys.exit(EXIT_CODE_NOK)

        file_handler = logging.FileHandler(config.log_filepath)
        stream_handler = logging.StreamHandler(sys.stderr)
        formatter = logging.Formatter("%(asctime)s - %(name)s - %(levelname)s - %(message)s")
        file_handler.setFormatter(formatter)
        stream_handler.setFormatter(formatter)
        handlers: list[logging.StreamHandler] = [file_handler] if quiet else [file_handler, stream_handler]
        logging.basicConfig(level=logging.INFO, handlers=handlers)
        start = datetime.now()

        actions_set = (
            {"build", "process-aux-dsets", "download", "convert", "analyze"} if "all" in actions else set(actions)
        )

        dset = build_or_load_dataset(framework, inputpath, "build" in actions_set, outputpath)

        processing_step: ProcessingStep
        for processing_step in [x for x in steps if x.name in actions_set]:
            processing_step.run(dset)

        end = datetime.now()
        logger.info(f"The computation took {(end - start)} seconds.")
    except Exception as e:
        click.echo(
            f"Unhandled exception: {e}",
            err=True,
        )
        return EXIT_CODE_NOK
    return EXIT_CODE_OK


if __name__ == "__main__":
    main()