service/src/piracyshield_service/forensic/create_archive.py

from __future__ import annotations

from piracyshield_service.base import BaseService

from piracyshield_component.config import Config
from piracyshield_component.utils.time import Time
from piracyshield_component.exception import ApplicationException

from piracyshield_data_model.forensic.archive.model import ForensicArchiveModel, ForensicArchiveModelNameException

from piracyshield_data_storage.forensic.storage import ForensicStorage, ForensicStorageCreateException, ForensicStorageGetException, ForensicStorageUpdateException

from piracyshield_data_storage.cache.storage import CacheStorage

from piracyshield_service.log.ticket.create import LogTicketCreateService

from piracyshield_service.forensic.tasks.analyze_forensic_archive import analyze_forensic_archive_task_caller

from piracyshield_service.forensic.errors import ForensicErrorCode, ForensicErrorMessage

import os

class ForensicCreateArchiveService(BaseService):

    """
    Manages the upload and schedules the analysis for the evidence archive.
    """

    importer_save_file_service = None

    log_ticket_create_service = None

    data_storage_cache = None

    data_storage = None

    data_model = None

    application_archive_config = None

    def __init__(self):
        """
        Inizialize logger and required modules.
        """

        super().__init__()

        self._prepare_configs()

        self._prepare_modules()

    def execute(self, ticket_id: str, archive_name: str) -> bool | Exception:
        """
        :param ticket_id: the ticket identifier related to the archive.
        :param archive_name: the name of the archive.
        :param archive_content: the data content of the file.
        :return: true if everything is successful.
        """

        model = self._validate_parameters(ticket_id, archive_name)

        # before getting to the expensive operations, let's check if the ticket exists
        if not self._ticket_id_exists(ticket_id):
            raise ApplicationException(ForensicErrorCode.NO_HASH_FOR_TICKET, ForensicErrorMessage.NO_HASH_FOR_TICKET)

        # the file is already in the cache
        if not self.data_storage_cache.exists(archive_name):
            self.logger.debug(f'Forensic evidence `{archive_name}` not found in cache')

            raise ApplicationException(ForensicErrorCode.GENERIC, ForensicErrorMessage.GENERIC)

        extension = self._get_extension(archive_name)

        if self._has_supported_extension(extension) == False:
            raise ApplicationException(ForensicErrorCode.EXTENSION_NOT_SUPPORTED, ForensicErrorMessage.EXTENSION_NOT_SUPPORTED)

        self.logger.debug(f'New forensic evidence `{archive_name}` in cache for ticket `{ticket_id}`')

        try:
            # update ticket with the archive name
            self.data_storage.update_archive_name(
                ticket_id = ticket_id,
                archive_name = archive_name,
                status = model.get('status'),
                updated_at = Time.now_iso8601()
            )

        except ForensicStorageUpdateException as e:
            self.logger.error(f"Could not update the ticket's forensic archive name")

            raise ApplicationException(ForensicErrorCode.GENERIC, ForensicErrorMessage.GENERIC, e)

        self._schedule_task(
            ticket_id = ticket_id
        )

        self.logger.info(f'Scheduled analysis for forensic evidence `{archive_name}` for ticket `{ticket_id}`')

        self.log_ticket_create_service.execute(
            ticket_id = ticket_id,
            message = f'Scheduled analysis for new forensic evidence package `{archive_name}`.'
        )

        return True

    def _ticket_id_exists(self, ticket_id: str) -> bool | Exception:
        """
        Check if the ticket exists.

        :param ticket_id: a ticket identifier.
        :return: true if everything is successful.
        """

        try:
            response = self.data_storage.exists_ticket_id(
                ticket_id = ticket_id
            )

            if response.next():
                self.logger.debug(f'Ticket found for `{ticket_id}`')

                return True

            return False

        except ForensicStorageGetException as e:
            self.logger.error(f'Could not verify `{ticket_id}` existence')

            raise ApplicationException(ForensicErrorCode.GENERIC, ForensicErrorMessage.GENERIC, e)

    def _get_extension(self, filename: str) -> str:
        """
        Extracts the extension from the filename.

        :param filename: the name of the file.
        :return: the extension of the file.
        """

        _, extension = os.path.splitext(filename)

        return extension

    def _has_supported_extension(self, extension: str) -> bool:
        """
        Checks wether the extension is supported.

        :param extension: the extension of the file.
        :return: true if supported.
        """

        return extension.lower() in self.application_archive_config.get('supported_extensions')

    def _schedule_task(self, ticket_id: str) -> bool | Exception:
        """
        Schedules the archive analysis.

        :param ticket_id: a ticket identifier.
        :return: true if everything is successful.
        """

        try:
            # schedule package analysis and upload to storage.
            analysis_task_id = self.task_service.create(
                task_caller = analyze_forensic_archive_task_caller,
                delay = 1,
                ticket_id = ticket_id
            )

            return True

        except Exception as e:
            self.logger.error(f'Could not create the task for `{ticket_id}`')

            raise ApplicationException(ForensicErrorCode.GENERIC, ForensicErrorMessage.GENERIC, e)

    def _validate_parameters(self, ticket_id: str, archive_name: str) -> bool | Exception:
        """
        Validates the inputs.

        :param ticket_id: a ticket identifier.
        :param archive_name: the name of the archive.
        :return: true if everything is successful.
        """

        try:
            model = self.data_model(
                ticket_id = ticket_id,
                name = archive_name
            )

            return model.to_dict()

        except ForensicArchiveModelNameException:
            raise ApplicationException(ForensicErrorCode.ARCHIVE_NAME, ForensicErrorMessage.ARCHIVE_NAME)

    def _prepare_configs(self) -> None:
        """
        Loads the configs.
        """

        self.application_archive_config = Config('application').get('archive')

    def _prepare_modules(self):
        """
        Initialize and set the instances.
        """

        self.data_model = ForensicArchiveModel

        self.data_storage = ForensicStorage()

        self.data_storage_cache = CacheStorage()

        self.log_ticket_create_service = LogTicketCreateService()