Fix: gnupg on encryption error (#117)

* catching gnupg error on file encryption. * added tests for parallel batch enc & dec * Update changelog * bump version v24.47.0 -> v24.47.1 * Ensure dummy key files exist so as not to break all tests * Fix dummy file creation * Alter PGP key file creation --------- Co-authored-by: Adam McDonagh <[email protected]>
adammcdonagh · Nov 22, 2024 · 7866b04 · 7866b04
1 parent c56accd
commit 7866b04
Show file tree

Hide file tree

Showing 14 changed files with 271 additions and 5 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -1,5 +1,9 @@
 # Changelog
 
+# v24.47.1
+
+- Add missing try/catch to decryption code when trying to delete gnupg directory
+
 # v24.47.0
 
 - Refactored dependency checker method to support Running/Failed state

diff --git a/pyproject.toml b/pyproject.toml
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "opentaskpy"
-version = "v24.47.0"
+version = "v24.47.1"
 authors = [{ name = "Adam McDonagh", email = "[email protected]" }]
 license = { text = "GPLv3" }
 classifiers = [
@@ -71,7 +71,7 @@ otf-batch-validator = "opentaskpy.cli.batch_validator:main"
 profile = 'black'
 
 [tool.bumpver]
-current_version = "v24.47.0"
+current_version = "v24.47.1"
 version_pattern = "vYY.WW.PATCH[-TAG]"
 commit_message = "bump version {old_version} -> {new_version}"
 commit = true

diff --git a/src/opentaskpy/taskhandlers/transfer.py b/src/opentaskpy/taskhandlers/transfer.py
@@ -772,7 +772,12 @@ def encrypt_files(
 
         # Remove the temporary gnupg keychain files under f"{tmpdir}/.gnupg"
         if path.exists(f"{tmpdir}/.gnupg"):
-            shutil.rmtree(f"{tmpdir}/.gnupg")
+            try:
+                shutil.rmtree(f"{tmpdir}/.gnupg")
+            except FileNotFoundError as e:
+                self.logger.warning(
+                    f".gnupg deletion failed - FileNotFound but continuing: {e}"
+                )
 
         return encrypted_files
 

diff --git a/test/cfg/transfers/decrypt-file-1-sftp.json.j2 b/test/cfg/transfers/decrypt-file-1-sftp.json.j2
@@ -0,0 +1,25 @@
+{
+  "type": "transfer",
+  "source": {
+    "hostname": "172.16.0.22",
+    "directory": "/home/application/testFiles/dest",
+    "fileRegex": "file-encrypt1\\.txt.gpg",
+    "protocol": { "name": "sftp", "credentials": { "username": "application" } }
+  },
+  "destination": [
+    {
+      "hostname": "172.16.0.21",
+      "directory": "/home/application/testFiles/dest",
+      "protocol": {
+        "name": "sftp",
+        "credentials": { "username": "application" },
+        "supportsPosixRename": false
+      },
+      "encryption": { "decrypt": true, "private_key": "{{ PRIVATE_KEY1 }}" },
+      "rename": {
+        "pattern": "^file-encrypt1\\.txt.gpg$",
+        "sub": "file-encrypt1-decrypted.txt"
+      }
+    }
+  ]
+}
diff --git a/test/cfg/transfers/decrypt-file-2-sftp.json.j2 b/test/cfg/transfers/decrypt-file-2-sftp.json.j2
@@ -0,0 +1,25 @@
+{
+  "type": "transfer",
+  "source": {
+    "hostname": "172.16.0.22",
+    "directory": "/home/application/testFiles/dest",
+    "fileRegex": "file-encrypt2\\.txt.gpg",
+    "protocol": { "name": "sftp", "credentials": { "username": "application" } }
+  },
+  "destination": [
+    {
+      "hostname": "172.16.0.21",
+      "directory": "/home/application/testFiles/dest",
+      "protocol": {
+        "name": "sftp",
+        "credentials": { "username": "application" },
+        "supportsPosixRename": false
+      },
+      "encryption": { "decrypt": true, "private_key": "{{ PRIVATE_KEY2 }}" },
+      "rename": {
+        "pattern": "^file-encrypt2\\.txt.gpg$",
+        "sub": "file-encrypt2-decrypted.txt"
+      }
+    }
+  ]
+}
diff --git a/test/cfg/transfers/encrypt-file-1-sftp.json.j2 b/test/cfg/transfers/encrypt-file-1-sftp.json.j2
@@ -0,0 +1,20 @@
+{
+  "type": "transfer",
+  "source": {
+    "hostname": "172.16.0.21",
+    "directory": "/home/application/testFiles/src",
+    "fileRegex": "file-encrypt1\\.txt",
+    "protocol": { "name": "sftp", "credentials": { "username": "application" } }
+  },
+  "destination": [
+    {
+      "hostname": "172.16.0.22",
+      "directory": "/home/application/testFiles/dest",
+      "protocol": {
+        "name": "sftp",
+        "credentials": { "username": "application" }
+      },
+      "encryption": { "encrypt": true, "public_key": "{{ PUBLIC_KEY1 }}" }
+    }
+  ]
+}
diff --git a/test/cfg/transfers/encrypt-file-2-sftp.json.j2 b/test/cfg/transfers/encrypt-file-2-sftp.json.j2
@@ -0,0 +1,20 @@
+{
+  "type": "transfer",
+  "source": {
+    "hostname": "172.16.0.21",
+    "directory": "/home/application/testFiles/src",
+    "fileRegex": "file-encrypt2\\.txt",
+    "protocol": { "name": "sftp", "credentials": { "username": "application" } }
+  },
+  "destination": [
+    {
+      "hostname": "172.16.0.22",
+      "directory": "/home/application/testFiles/dest",
+      "protocol": {
+        "name": "sftp",
+        "credentials": { "username": "application" }
+      },
+      "encryption": { "encrypt": true, "public_key": "{{ PUBLIC_KEY2 }}" }
+    }
+  ]
+}
diff --git a/test/cfg/variables.json.j2 b/test/cfg/variables.json.j2
@@ -26,5 +26,9 @@
   "testLookup": "{{ lookup('file', path='/tmp/variable_lookup.txt') }}",
   "testLookup2": "{{ lookup('http_json', url='https://jsonplaceholder.typicode.com/posts/1', jsonpath='$.title') }}",
   "global_protocol_vars": {"name": "email", "smtp_port": 587, "smtp_server": "smtp.gmail.com"},
-  "SSH_VARS": {"name": "ssh", "credentials": {"username": "someuser"}}
+  "SSH_VARS": {"name": "ssh", "credentials": {"username": "someuser"}},
+  "PUBLIC_KEY1": "{{ lookup('file', path='/tmp/public_key_1.txt') }}",
+  "PUBLIC_KEY2": "{{ lookup('file', path='/tmp/public_key_2.txt') }}",
+  "PRIVATE_KEY1": "{{ lookup('file', path='/tmp/private_key_1.txt') }}",
+  "PRIVATE_KEY2": "{{ lookup('file', path='/tmp/private_key_2.txt') }}"
 }
diff --git a/test/createTestFiles.sh b/test/createTestFiles.sh
@@ -2,6 +2,10 @@
 
 # Create dummy variable for lookup command
 echo "file_variable" > /tmp/variable_lookup.txt
+echo "" > /tmp/private_key_1.txt
+echo "" > /tmp/private_key_2.txt
+echo "" > /tmp/public_key_1.txt
+echo "" > /tmp/public_key_2.txt
 
 # Get current script directory
 DIR="$( cd "$( dirname "${BASH_SOURCE[0]}" )" && pwd )"

diff --git a/tests/fixtures/pgp.py b/tests/fixtures/pgp.py
@@ -1,5 +1,8 @@
 # pylint: skip-file
+import os
+
 import pytest
+from pytest_shell import fs
 
 
 @pytest.fixture(scope="session")
@@ -262,3 +265,31 @@ def private_key_2() -> str:
 5FJn9I4DYA==
 =cmHi
 -----END PGP PRIVATE KEY BLOCK-----"""
+
+
+@pytest.fixture(scope="function")
+def store_pgp_keys(public_key, public_key_2, private_key, private_key_2) -> bool:
+    pub1 = "\\\\n".join(public_key.splitlines())
+    pub2 = "\\\\n".join(public_key_2.splitlines())
+    prv1 = "\\\\n".join(private_key.splitlines())
+    prv2 = "\\\\n".join(private_key_2.splitlines())
+
+    # Remove all the files and the recreate them
+    if os.path.exists("/tmp/public_key_1.txt"):
+        os.remove("/tmp/public_key_1.txt")
+    if os.path.exists("/tmp/public_key_2.txt"):
+        os.remove("/tmp/public_key_2.txt")
+    if os.path.exists("/tmp/private_key_1.txt"):
+        os.remove("/tmp/private_key_1.txt")
+    if os.path.exists("/tmp/private_key_2.txt"):
+        os.remove("/tmp/private_key_2.txt")
+
+    fs.create_files(
+        [
+            {"/tmp/public_key_1.txt": {"content": pub1}},
+            {"/tmp/public_key_2.txt": {"content": pub2}},
+            {"/tmp/private_key_1.txt": {"content": prv1}},
+            {"/tmp/private_key_2.txt": {"content": prv2}},
+        ]
+    )
+    return True
diff --git a/tests/fixtures/ssh_clients.py b/tests/fixtures/ssh_clients.py
@@ -22,7 +22,15 @@ def env_vars() -> None:
         del os.environ["OTF_NO_LOG"]
 
     # We're using the proper config file for this, so we need to make sure something exist in /tmp/variable_lookup.txt
-    fs.create_files([{"/tmp/variable_lookup.txt": {"content": "test1234"}}])
+    fs.create_files(
+        [
+            {"/tmp/variable_lookup.txt": {"content": "test1234"}},
+            {"/tmp/public_key_1.txt": {"content": "test1234"}},
+            {"/tmp/public_key_2.txt": {"content": "test1234"}},
+            {"/tmp/private_key_1.txt": {"content": "test1234"}},
+            {"/tmp/private_key_2.txt": {"content": "test1234"}},
+        ]
+    )
 
 
 @pytest.fixture(scope="session")

diff --git a/tests/test_config_loader.py b/tests/test_config_loader.py
@@ -47,6 +47,10 @@ def test_load_variables(tmpdir):
                     "content": "hello",
                 }
             },
+            {"/tmp/public_key_1.txt": {"content": "test1234"}},
+            {"/tmp/public_key_2.txt": {"content": "test1234"}},
+            {"/tmp/private_key_1.txt": {"content": "test1234"}},
+            {"/tmp/private_key_2.txt": {"content": "test1234"}},
         ]
     )
 

diff --git a/tests/test_docker_task_run.py b/tests/test_docker_task_run.py
@@ -203,6 +203,14 @@ def test_docker_run(
         f"{log_dir}:/logs",
         "--volume",
         "/tmp/variable_lookup.txt:/tmp/variable_lookup.txt",
+        "--volume",
+        "/tmp/public_key_1.txt:/tmp/public_key_1.txt",
+        "--volume",
+        "/tmp/public_key_2.txt:/tmp/public_key_2.txt",
+        "--volume",
+        "/tmp/private_key_1.txt:/tmp/private_key_1.txt",
+        "--volume",
+        "/tmp/private_key_2.txt:/tmp/private_key_2.txt",
         "-e",
         "OTF_SSH_KEY=/test/testFiles/id_rsa",
         image_name_dev,
@@ -258,6 +266,14 @@ def test_standard_docker_image(
         "--volume",
         "/tmp/variable_lookup.txt:/tmp/variable_lookup.txt",
         "--volume",
+        "/tmp/public_key_1.txt:/tmp/public_key_1.txt",
+        "--volume",
+        "/tmp/public_key_2.txt:/tmp/public_key_2.txt",
+        "--volume",
+        "/tmp/private_key_1.txt:/tmp/private_key_1.txt",
+        "--volume",
+        "/tmp/private_key_2.txt:/tmp/private_key_2.txt",
+        "--volume",
         f"{root_dir}/testLogs:/logs",
         "--env",
         "OTF_RUN_ID=docker_log_test",

diff --git a/tests/test_taskhandler_batch_encryption.py b/tests/test_taskhandler_batch_encryption.py
@@ -0,0 +1,100 @@
+# pylint: skip-file
+# ruff: noqa
+import os
+
+from pytest_shell import fs
+
+import opentaskpy.otflogging
+from opentaskpy.config.loader import ConfigLoader
+from opentaskpy.taskhandlers import batch, execution, transfer
+from tests.file_helper import *  # noqa: F403
+from tests.fixtures.pgp import *  # noqa: F403, F401
+from tests.fixtures.ssh_clients import *  # noqa: F403
+
+os.environ["OTF_NO_LOG"] = "1"
+os.environ["OTF_LOG_LEVEL"] = "DEBUG"
+
+parallel_encryption_batch = {
+    "type": "batch",
+    "tasks": [
+        {"order_id": 1, "task_id": "encrypt-file-1-sftp"},
+        {"order_id": 2, "task_id": "encrypt-file-2-sftp"},
+    ],
+}
+
+parallel_decryption_batch = {
+    "type": "batch",
+    "tasks": [
+        {"order_id": 1, "task_id": "decrypt-file-1-sftp"},
+        {"order_id": 2, "task_id": "decrypt-file-2-sftp"},
+    ],
+}
+
+
+def test_batch_parallel_encryption(
+    root_dir, env_vars, clear_logs, store_pgp_keys, setup_sftp_keys
+):
+    # Assert one of the keys randomly
+    assert os.path.exists("/tmp/public_key_1.txt")
+
+    config_loader = ConfigLoader("test/cfg")
+
+    fs.create_files(
+        [
+            {
+                f"{root_dir}/testFiles/sftp_1/src/file-encrypt1.txt": {
+                    "content": "hellothere!"
+                }
+            },
+            {
+                f"{root_dir}/testFiles/sftp_1/src/file-encrypt2.txt": {
+                    "content": "hellothere!"
+                }
+            },
+        ]
+    )
+    assert os.path.exists(f"{root_dir}/testFiles/sftp_1/src/file-encrypt1.txt")
+    assert os.path.exists(f"{root_dir}/testFiles/sftp_1/src/file-encrypt2.txt")
+
+    batch_obj = batch.Batch(
+        None,
+        f"parallel-encryption-batch",
+        parallel_encryption_batch,
+        config_loader,
+    )
+
+    # Run the batch and expect a true status
+    assert batch_obj.run()
+
+    assert os.path.exists(f"{root_dir}/testFiles/sftp_2/dest/file-encrypt1.txt.gpg")
+    assert os.path.exists(f"{root_dir}/testFiles/sftp_2/dest/file-encrypt2.txt.gpg")
+
+
+def test_batch_parallel_decryption(
+    root_dir, env_vars, clear_logs, store_pgp_keys, setup_sftp_keys
+):
+    # Assert one of the keys randomly
+    assert os.path.exists("/tmp/private_key_1.txt")
+
+    config_loader = ConfigLoader("test/cfg")
+
+    assert os.path.exists(f"{root_dir}/testFiles/sftp_2/dest/file-encrypt1.txt.gpg")
+    assert os.path.exists(f"{root_dir}/testFiles/sftp_2/dest/file-encrypt2.txt.gpg")
+
+    batch_obj = batch.Batch(
+        None,
+        f"parallel-decryption-batch",
+        parallel_decryption_batch,
+        config_loader,
+    )
+
+    # Run the batch and expect a true status
+    assert batch_obj.run()
+
+    # Assert files have been decrypted and transferred back to sftp1
+    assert os.path.exists(
+        f"{root_dir}/testFiles/sftp_1/dest/file-encrypt1-decrypted.txt"
+    )
+    assert os.path.exists(
+        f"{root_dir}/testFiles/sftp_1/dest/file-encrypt2-decrypted.txt"
+    )