Enh/custom pkl (#6)

* add functionality for custom pkl file

* add functionality for custom pkl files and zipped archives

osp/models/catalytic/co_catalyticfoam.py

@@ -248,6 +248,12 @@ def __post_init_post_parse__(self):
                 calc.add(self._make_species(species), rel=emmo.hasCalculationInput)
             calc.add(self._make_continuum_model(), rel=emmo.hasCalculationInput)
             calc.add(self._make_catalyist_amount(), rel=emmo.hasCalculationInput)
+            if self.species_from_upload:
+                pkl = emmo.PKLFile(uid=self.species_from_upload)
+                calc.add(pkl, rel=emmo.hasCalculationInput)
+            if self.patches_from_upload:
+                tar = emmo.TarballFile(uid=self.patches_from_upload)
+                calc.add(tar, rel=emmo.hasCalculationInput)
         file = tempfile.NamedTemporaryFile(suffix=".ttl")
         export_cuds(session, file.name)
         self._uuid = get_upload(file)

osp/wrappers/simcatalyticfoam/catalyticfoam_engine.py

@@ -2,8 +2,9 @@
 import os
 import shutil
 import subprocess  # nosec
+import tarfile
 import tempfile
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Optional
 
 import osp.dictionaries.catalyticFoam as case
 
@@ -37,6 +38,8 @@ def __init__(self, case: str, config: dict = {}):
                 directory=runtime,
                 filename=tar,
             )
+        self._pkl: Optional[str] = None
+        self._input_tarball: Optional[str] = None
         self._parse_files = dict()
         self._config: dict = config
         self._current_process: subprocess.Popen = None
@@ -66,6 +69,12 @@ def run(self) -> None:
             with open(path, "w+") as file:
                 file_content = serialize(content)
                 file.write(file_content)
+        if self._pkl:
+            logger.info("Will use custom PKL for catalytic wall: `%s`", self._pkl)
+            dest_pkl = os.path.join(
+                self._config["directory"], "ml_ExtraTrees_forCFD.pkl"
+            )
+            shutil.copy(self._pkl, dest_pkl)
         for command in self._config["commands"]:
             cmd = self._prepare_command(command)
             logger.info(
@@ -221,3 +230,42 @@ def exit_code(cls):
     @property
     def tarball(cls):
         return f"{cls._config['filename']}.tar"
+
+    @property
+    def input_tarball(cls):
+        return cls._input_tarball
+
+    @input_tarball.setter
+    def _set_tarball(self, value):
+        self._input_tarball = value
+        logger.info("Received new value for case directory to be used: `%s`.", value)
+        if "gz" in value:
+            mode = "r:gz"
+        elif "bz2" in value:
+            mode = "r:bz2"
+        else:
+            mode = "r"
+        logger.info(
+            "Will remove previous files in runtime directory: `%s`",
+            self._config["directory"],
+        )
+        shutil.rmtree(self._config["directory"])
+        logger.info(
+            "Will untar provided file in to runtime directory: `%s`",
+            self._config["directory"],
+        )
+        allowed_members = settings.allowed_members.strip().split(";")
+        with tarfile.open(value, mode=mode) as file:
+            # Iterate over the members and validate
+            for member in file.getmembers():
+                if member.name in allowed_members:
+                    file.extract(
+                        member, path=self._config["directory"], set_attrs=False
+                    )
+                else:
+                    logger.warning("Skipping member: %s", member.name)
+        logger.info("Untaring complete.")
+
+    @property
+    def pkl(cls):
+        return cls._pkl

osp/wrappers/simcatalyticfoam/settings.py

@@ -53,6 +53,14 @@ class CatalyticFoamSettings(BaseSettings):
         E.g. `fixedValue` for a Dirichlet boundary condition.""",
     )
 
+    allowed_members: str = Field(
+        "0;constant;kinetic;system;CatalyticReactors;ml_ExtraTrees_forCFD.pkl",
+        description="""If an tar-file is used as custom use case directory, the
+                                  `;`-seprated values in the provided string are the only members which
+                                  are allowed to be extracted from the tarball-file (due to security reasons,
+                                  see [B202:tarfile_unsafe_members] for reference.)""",
+    )
+
 
 class ReaxProSettings(ModelSettings, CatalyticFoamSettings):
     """General Reaxpro wrapper settings"""

osp/wrappers/simcatalyticfoam/simcatalyticfoam.py

@@ -5,6 +5,7 @@
 from osp.core.namespaces import emmo
 from osp.core.session import SimWrapperSession
 from osp.models.catalytic.utils import check_arcp, wrap_arcp
+from osp.models.utils.general import get_download
 
 from .catalyticfoam_engine import CatalyticFoamEngine, settings
 
@@ -34,6 +35,9 @@ def _run(self, root_object) -> None:
 
     # OVERRIDE
     def _apply_added(self, root_object, buffer) -> None:
+        for obj in buffer.values():
+            if obj.is_a(emmo.TarballFile):
+                self._engine.input_tarball = get_download(str(obj.uid), as_file=True)
         for obj in buffer.values():
             self._wrap(obj)
 
@@ -68,6 +72,8 @@ def _wrap(self, cuds: "Cuds") -> None:
             self._check_for_value(cuds)
         elif cuds.is_a(emmo.AdjustableSimulationTimeStep):
             self._wrap_arcp(cuds, "yes")
+        elif cuds.is_a(emmo.PKLFile):
+            self._engine.pkl = get_download(str(cuds.uid), as_file=True)
 
     def _make_new_chemical(self, composition: "Cuds") -> None:
         species = composition.get(oclass=emmo.ChemicalSpecies)