exporters: support reading keys from env vars (fixes #97)

murakami/exporters/gcs.py

@@ -1,7 +1,10 @@
-import os
+import base64
 import io
-import subprocess
+import json
 import logging
+import os
+import subprocess
+
 import jsonlines
 
 from google.cloud import storage
@@ -33,7 +36,13 @@ def __init__(
         self.key = config.get("key", None)
         # Initialize a GCS Client object from the provided key.
         # This client will be reused for all the subsequent GCS uploads.
-        self.client = storage.Client.from_service_account_json(self.key)
+        key_content = os.environ.get("MURAKAMI_GCS_KEY_CONTENT", None)
+        if key_content is not None:
+            logger.debug("GCS: loading credentials from MURAKAMI_GCS_KEY_CONTENT")
+            key_dict = json.loads(base64.b64decode(key_content).decode("utf-8"))
+            self.client = storage.Client.from_service_account_info(key_dict)
+        else:
+            self.client = storage.Client.from_service_account_json(self.key)
 
     def upload(self, data, bucket_name, object_name):
         if self.client is None:

murakami/exporters/scp.py

@@ -1,6 +1,8 @@
+import base64
 import io
 import logging
 import os
+import tempfile
 
 import jsonlines
 from paramiko import SSHClient
@@ -38,6 +40,7 @@ def __init__(
         self.username = config.get("username", None)
         self.password = config.get("password", None)
         self.private_key = config.get("key", None)
+        self._key_content = os.environ.get("MURAKAMI_SCP_KEY_CONTENT", None)
 
     def _push_single(self, test_name="", data=None, timestamp=None,
         test_idx=None):
@@ -46,7 +49,7 @@ def _push_single(self, test_name="", data=None, timestamp=None,
             logger.error("scp.target must be specified")
             return
 
-        if self.username is None and self.private_key is None:
+        if self.username is None and self.private_key is None and self._key_content is None:
             logging.error("scp.username or scp.private_key must be provided.")
 
         try:
@@ -58,14 +61,23 @@ def _push_single(self, test_name="", data=None, timestamp=None,
         ssh = SSHClient()
         ssh.set_missing_host_key_policy(AutoAddPolicy)
 
+        tmp_key_file = None
         try:
+            key_filename = self.private_key
+            if self._key_content is not None:
+                logger.debug("SCP: loading key from MURAKAMI_SCP_KEY_CONTENT")
+                tmp_key_file = tempfile.NamedTemporaryFile(delete=False)
+                tmp_key_file.write(base64.b64decode(self._key_content))
+                tmp_key_file.flush()
+                tmp_key_file.close()
+                key_filename = tmp_key_file.name
             ssh.connect(
                 dst_host,
                 int(self.port),
                 username=self.username,
                 password=self.password,
                 timeout=defaults.SSH_TIMEOUT,
-                key_filename=self.private_key,
+                key_filename=key_filename,
             )
 
             with SCPClient(ssh.get_transport()) as scp:
@@ -79,3 +91,5 @@ def _push_single(self, test_name="", data=None, timestamp=None,
             logger.error("SCP exporter failed: %s", err)
         finally:
             ssh.close()
+            if tmp_key_file is not None:
+                os.unlink(tmp_key_file.name)

tests/test_gcs_exporter.py

@@ -0,0 +1,57 @@
+import base64
+import json
+import os
+from unittest.mock import MagicMock, patch
+
+
+FAKE_KEY_DICT = {
+    "type": "service_account",
+    "project_id": "test-project",
+    "private_key_id": "key-id",
+    "private_key": "-----BEGIN RSA PRIVATE KEY-----\nMIIEowIBAAKCAQEA\n-----END RSA PRIVATE KEY-----\n",
+    "client_email": "test@test-project.iam.gserviceaccount.com",
+    "client_id": "123456789",
+    "auth_uri": "https://accounts.google.com/o/oauth2/auth",
+    "token_uri": "https://oauth2.googleapis.com/token",
+}
+
+FAKE_KEY_CONTENT = base64.b64encode(
+    json.dumps(FAKE_KEY_DICT).encode("utf-8")
+).decode("utf-8")
+
+
+@patch("murakami.exporters.gcs.storage.Client")
+def test_gcs_env_var_uses_from_service_account_info(mock_client_cls):
+    mock_client_cls.from_service_account_info.return_value = MagicMock()
+
+    env = {"MURAKAMI_GCS_KEY_CONTENT": FAKE_KEY_CONTENT}
+    with patch.dict(os.environ, env, clear=False):
+        from murakami.exporters.gcs import GCSExporter
+        exporter = GCSExporter(
+            name="test",
+            config={"target": "gs://bucket/path"},
+        )
+
+    mock_client_cls.from_service_account_info.assert_called_once_with(
+        FAKE_KEY_DICT
+    )
+    mock_client_cls.from_service_account_json.assert_not_called()
+
+
+@patch("murakami.exporters.gcs.storage.Client")
+def test_gcs_file_path_used_when_no_env_var(mock_client_cls):
+    mock_client_cls.from_service_account_json.return_value = MagicMock()
+
+    env_without_var = {k: v for k, v in os.environ.items()
+                       if k != "MURAKAMI_GCS_KEY_CONTENT"}
+    with patch.dict(os.environ, env_without_var, clear=True):
+        from murakami.exporters.gcs import GCSExporter
+        exporter = GCSExporter(
+            name="test",
+            config={"target": "gs://bucket/path", "key": "/path/to/key.json"},
+        )
+
+    mock_client_cls.from_service_account_json.assert_called_once_with(
+        "/path/to/key.json"
+    )
+    mock_client_cls.from_service_account_info.assert_not_called()

tests/test_scp_exporter.py

@@ -0,0 +1,67 @@
+import base64
+import os
+from unittest.mock import MagicMock, call, patch
+
+
+FAKE_PEM = b"-----BEGIN RSA PRIVATE KEY-----\nMIIEowIBAAKCAQEA\n-----END RSA PRIVATE KEY-----\n"
+FAKE_KEY_CONTENT = base64.b64encode(FAKE_PEM).decode("utf-8")
+
+
+@patch("murakami.exporters.scp.SCPClient")
+@patch("murakami.exporters.scp.SSHClient")
+def test_scp_env_var_writes_temp_file(mock_ssh_cls, mock_scp_cls):
+    mock_ssh = MagicMock()
+    mock_ssh_cls.return_value = mock_ssh
+    mock_ssh.get_transport.return_value = MagicMock()
+
+    mock_scp_instance = MagicMock()
+    mock_scp_cls.return_value.__enter__ = MagicMock(return_value=mock_scp_instance)
+    mock_scp_cls.return_value.__exit__ = MagicMock(return_value=False)
+
+    env = {"MURAKAMI_SCP_KEY_CONTENT": FAKE_KEY_CONTENT}
+    with patch.dict(os.environ, env, clear=False):
+        from murakami.exporters.scp import SCPExporter
+        exporter = SCPExporter(
+            name="test",
+            config={
+                "target": "host:/remote/path",
+                "username": "user",
+            },
+        )
+        exporter._push_single(test_name="ndt7", data='{"result": 1}',
+                               timestamp="2024-01-01T00:00:00.000000")
+
+    connect_kwargs = mock_ssh.connect.call_args
+    key_filename = connect_kwargs[1]["key_filename"]
+    assert key_filename is not None
+    assert not os.path.exists(key_filename), "temp file must be deleted after push"
+
+
+@patch("murakami.exporters.scp.SCPClient")
+@patch("murakami.exporters.scp.SSHClient")
+def test_scp_file_path_used_when_no_env_var(mock_ssh_cls, mock_scp_cls):
+    mock_ssh = MagicMock()
+    mock_ssh_cls.return_value = mock_ssh
+    mock_ssh.get_transport.return_value = MagicMock()
+
+    mock_scp_instance = MagicMock()
+    mock_scp_cls.return_value.__enter__ = MagicMock(return_value=mock_scp_instance)
+    mock_scp_cls.return_value.__exit__ = MagicMock(return_value=False)
+
+    env_without_var = {k: v for k, v in os.environ.items()
+                       if k != "MURAKAMI_SCP_KEY_CONTENT"}
+    with patch.dict(os.environ, env_without_var, clear=True):
+        from murakami.exporters.scp import SCPExporter
+        exporter = SCPExporter(
+            name="test",
+            config={
+                "target": "host:/remote/path",
+                "username": "user",
+                "key": "/path/to/id_rsa",
+            },
+        )
+        exporter._push_single(test_name="ndt7", data='{"result": 1}',
+                               timestamp="2024-01-01T00:00:00.000000")
+
+    connect_kwargs = mock_ssh.connect.call_args
+    assert connect_kwargs[1]["key_filename"] == "/path/to/id_rsa"

utilities/encode_key.sh

@@ -0,0 +1,18 @@
+#!/bin/sh
+# Usage: encode_key.sh /path/to/key.json   (GCS service account JSON)
+#        encode_key.sh /path/to/id_rsa      (SCP PEM private key)
+#
+# The encoded value can then be exported as:
+#   export MURAKAMI_GCS_KEY_CONTENT="<output>"   # for the GCS exporter
+#   export MURAKAMI_SCP_KEY_CONTENT="<output>"   # for the SCP exporter
+
+if [ -z "$1" ]; then
+    echo "Usage: $0 /path/to/keyfile" >&2
+    exit 1
+fi
+
+encoded=$(base64 -w 0 "$1")
+echo "$encoded"
+echo ""
+echo "Set the above value as MURAKAMI_GCS_KEY_CONTENT or MURAKAMI_SCP_KEY_CONTENT"
+echo "in your environment or container configuration."