rsa_keys: refactor

Author: pawbu

job_executor/adapter/fs/__init__.py

@@ -4,7 +4,9 @@
 
 from job_executor.adapter.fs.datastore_files import DatastoreDirectory
 from job_executor.adapter.fs.input_files import InputDirectory
+from job_executor.adapter.fs.private_keys_directory import PrivateKeysDirectory
 from job_executor.adapter.fs.working_files import WorkingDirectory
+from job_executor.config import environment
 
 
 class FileSystemAdapter(Protocol):
@@ -21,13 +23,17 @@ class LocalStorageAdapter:
     datastore_dir: DatastoreDirectory
     working_dir: WorkingDirectory
     input_dir: InputDirectory
+    private_keys_dir: PrivateKeysDirectory
 
-    def __init__(self, datastore_dir_path: Path) -> None:
+    def __init__(self, datastore_dir_path: Path, datastore_rdn: str) -> None:
         self.datastore_dir = DatastoreDirectory(datastore_dir_path)
         self.working_dir = WorkingDirectory(
             Path(f"{datastore_dir_path}_working")
         )
         self.input_dir = InputDirectory(Path(f"{datastore_dir_path}_input"))
+        self.private_keys_dir = PrivateKeysDirectory(
+            Path(environment.private_keys_dir) / datastore_rdn
+        )
 
     def move_working_dir_parquet_to_datastore(self, dataset_name: str) -> None:
         """

job_executor/adapter/fs/private_keys_directory.py

@@ -0,0 +1,27 @@
+import os
+from dataclasses import dataclass
+from pathlib import Path
+
+
+@dataclass
+class PrivateKeysDirectory:
+    path_with_rdn: Path
+
+    def create(self) -> bool:
+        if not self.path_with_rdn.exists():
+            os.makedirs(self.path_with_rdn)
+            return True
+        return False
+
+    def save_private_key(self, microdata_private_key_pem: bytes) -> None:
+        with open(self._get_private_key_location(), "wb") as file:
+            file.write(microdata_private_key_pem)
+
+    def clean_up(self) -> bool:
+        if self._get_private_key_location().exists():
+            os.remove(self._get_private_key_location())
+            return True
+        return False
+
+    def _get_private_key_location(self) -> Path:
+        return self.path_with_rdn / "microdata_private_key.pem"

job_executor/app.py

@@ -24,7 +24,7 @@ def initialize_app() -> Manager:
         rollback.fix_interrupted_jobs()
         for rdn in datastore_api.get_datastores():
             local_storage = LocalStorageAdapter(
-                datastore_api.get_datastore_directory(rdn)
+                datastore_api.get_datastore_directory(rdn), rdn
             )
             if local_storage.datastore_dir.temporary_backup_exists():
                 raise StartupException(f"tmp directory exists for {rdn}")

job_executor/domain/datastores.py

@@ -1,10 +1,4 @@
 import logging
-import os
-from pathlib import Path
-
-from cryptography.hazmat.backends import default_backend
-from cryptography.hazmat.primitives import serialization
-from cryptography.hazmat.primitives.asymmetric import rsa
 
 from job_executor.adapter import datastore_api
 from job_executor.adapter.datastore_api.models import JobStatus
@@ -25,12 +19,12 @@
     UnnecessaryUpdateException,
     VersioningException,
 )
-from job_executor.config import environment
 from job_executor.domain.models import JobContext
 from job_executor.domain.rollback import (
     rollback_bump,
     rollback_manager_phase_import_job,
 )
+from job_executor.domain.rsa_keys import generate_rsa_key_pair
 
 logger = logging.getLogger()
 
@@ -605,55 +599,34 @@ def generate_rsa_keys(
     """
     job_id = job_context.job.job_id
     datastore_rdn = job_context.job.datastore_rdn
+    private_keys_dir = job_context.local_storage.private_keys_dir
     try:
         logger.info(f"{job_id}: initiated")
         datastore_api.update_job_status(job_id, JobStatus.INITIATED)
 
-        target_dir = Path(environment.private_keys_dir) / datastore_rdn
-
-        if not target_dir.exists():
-            logger.info(
-                f"{job_id}: Creating private keys directory at {target_dir}"
-            )
-            os.makedirs(target_dir)
-
-        logger.info(f"{job_id}: Generating RSA key pair")
-        private_key = rsa.generate_private_key(
-            public_exponent=65537, key_size=2048, backend=default_backend()
-        )
-        public_key = private_key.public_key()
-
-        microdata_private_key_pem = private_key.private_bytes(
-            encoding=serialization.Encoding.PEM,
-            format=serialization.PrivateFormat.PKCS8,
-            encryption_algorithm=serialization.NoEncryption(),
+        logger.info(
+            f"{job_id}: Checking private keys directory at "
+            f"{private_keys_dir.path_with_rdn}"
         )
+        if private_keys_dir.create():
+            logger.info(f"{job_id}: Private keys directory created")
 
-        private_key_location = target_dir / "microdata_private_key.pem"
-        with open(private_key_location, "wb") as file:
-            file.write(microdata_private_key_pem)
-        logger.info(f"{job_id}: Saved private key to {private_key_location}")
+        logger.info(f"{job_id}: Generating RSA key pair")
+        private_key_pem, public_key_pem = generate_rsa_key_pair()
 
-        microdata_public_key_pem = public_key.public_bytes(
-            encoding=serialization.Encoding.PEM,
-            format=serialization.PublicFormat.SubjectPublicKeyInfo,
-        )
+        private_keys_dir.save_private_key(private_key_pem)
+        logger.info(f"{job_id}: Saved private key")
 
         try:
             logger.info(f"{job_id}: Posting public key to datastore-api")
-            datastore_api.post_public_key(
-                datastore_rdn, microdata_public_key_pem
-            )
+            datastore_api.post_public_key(datastore_rdn, public_key_pem)
         except Exception as post_error:
             logger.error(
                 f"{job_id}: Failed to post public key to datastore-api, "
                 "cleaning up saved private key"
             )
-            if private_key_location.exists():
-                os.remove(private_key_location)
-                logger.info(
-                    f"{job_id}: Deleted private key at {private_key_location}"
-                )
+            if private_keys_dir.clean_up():
+                logger.info(f"{job_id}: Deleted private key due to error")
             raise post_error
 
         logger.info(f"{job_id}: completed")

job_executor/domain/models.py

@@ -18,7 +18,8 @@ class JobContext:
 
 def build_job_context(job: Job, handler: handler_type) -> JobContext:
     local_storage = LocalStorageAdapter(
-        datastore_api.get_datastore_directory(job.datastore_rdn)
+        datastore_api.get_datastore_directory(job.datastore_rdn),
+        job.datastore_rdn,
     )
     job_size = (
         local_storage.input_dir.get_importable_tar_size_in_bytes(

job_executor/domain/rollback.py

@@ -28,7 +28,8 @@
 def rollback_bump(job: Job, bump_manifesto: DatastoreVersion) -> None:
     job_id = job.job_id
     local_storage = LocalStorageAdapter(
-        datastore_api.get_datastore_directory(job.datastore_rdn)
+        datastore_api.get_datastore_directory(job.datastore_rdn),
+        job.datastore_rdn,
     )
     try:
         logger.info(f"{job_id}: Restoring files from temporary backup")
@@ -138,7 +139,8 @@ def rollback_worker_phase_import_job(
 ) -> None:
     job_id = job.job_id
     local_storage = LocalStorageAdapter(
-        datastore_api.get_datastore_directory(job.datastore_rdn)
+        datastore_api.get_datastore_directory(job.datastore_rdn),
+        job.datastore_rdn,
     )
     logger.warning(
         f"{job_id}: Rolling back worker job "
@@ -194,7 +196,8 @@ def rollback_manager_phase_import_job(
     """
     job_id = job.job_id
     local_storage = LocalStorageAdapter(
-        datastore_api.get_datastore_directory(job.datastore_rdn)
+        datastore_api.get_datastore_directory(job.datastore_rdn),
+        job.datastore_rdn,
     )
     logger.warning(
         f"{job_id}: Rolling back import job "

job_executor/domain/rsa_keys.py

@@ -0,0 +1,27 @@
+import logging
+
+from cryptography.hazmat.backends import default_backend
+from cryptography.hazmat.primitives import serialization
+from cryptography.hazmat.primitives.asymmetric import rsa
+
+logger = logging.getLogger()
+
+
+def generate_rsa_key_pair() -> tuple[bytes, bytes]:
+    private_key = rsa.generate_private_key(
+        public_exponent=65537, key_size=2048, backend=default_backend()
+    )
+    public_key = private_key.public_key()
+
+    private_key_pem = private_key.private_bytes(
+        encoding=serialization.Encoding.PEM,
+        format=serialization.PrivateFormat.PKCS8,
+        encryption_algorithm=serialization.NoEncryption(),
+    )
+
+    public_key_pem = public_key.public_bytes(
+        encoding=serialization.Encoding.PEM,
+        format=serialization.PublicFormat.SubjectPublicKeyInfo,
+    )
+
+    return private_key_pem, public_key_pem

tests/integration/test_datastore.py

@@ -112,7 +112,7 @@ def generate_job_parameters(
 
     return JobContext(
         handler="worker",
-        local_storage=LocalStorageAdapter(DATASTORE_DIR),
+        local_storage=LocalStorageAdapter(DATASTORE_DIR, "TEST_DATASTORE"),
         job_size=100,
         job=Job(
             job_id="1",

tests/integration/test_import.py

@@ -88,7 +88,7 @@ def set_up_resources():
 def generate_job_context(operation: Operation, target: str) -> JobContext:
     return JobContext(
         handler="worker",
-        local_storage=LocalStorageAdapter(DATASTORE_DIR),
+        local_storage=LocalStorageAdapter(DATASTORE_DIR, "TEST_DATASTORE"),
         job_size=100,
         job=Job(
             job_id="1",

tests/unit/adapter/fs/model/test_datastore_versions.py

@@ -16,7 +16,7 @@ def load_json(file_path):
 
 DATASTORE_DIR = "tests/unit/resources/adapter/fs/TEST_DATASTORE"
 METADATA_DIR = f"{DATASTORE_DIR}/datastore"
-local_storage = LocalStorageAdapter(Path(DATASTORE_DIR))
+local_storage = LocalStorageAdapter(Path(DATASTORE_DIR), "TEST_DATASTORE")
 DATASTORE_VERSIONS_PATH = f"{METADATA_DIR}/datastore_versions.json"