diff --git a/murakami/exporters/gcs.py b/murakami/exporters/gcs.py index 74db067..6a62dc2 100644 --- a/murakami/exporters/gcs.py +++ b/murakami/exporters/gcs.py @@ -1,7 +1,9 @@ -import os +import base64 import io -import subprocess +import json import logging +import subprocess + import jsonlines from google.cloud import storage @@ -31,9 +33,23 @@ def __init__( logging.debug(config) self.target = config.get("target", None) self.key = config.get("key", None) + self.key_content = config.get("key_content", 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) + if self.key_content is not None: + logger.debug("GCS: loading credentials from key_content config") + try: + key_dict = json.loads( + base64.b64decode(self.key_content).decode("utf-8") + ) + self.client = storage.Client.from_service_account_info(key_dict) + except Exception as e: + logger.error( + "GCS: failed to load credentials from key_content: %s", e + ) + self.client = None + else: + self.client = storage.Client.from_service_account_json(self.key) def upload(self, data, bucket_name, object_name): if self.client is None: diff --git a/murakami/exporters/scp.py b/murakami/exporters/scp.py index 973a07a..824348d 100644 --- a/murakami/exporters/scp.py +++ b/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 = config.get("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,24 @@ 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 key_content config") + tmp_key_file = tempfile.NamedTemporaryFile(delete=False) + tmp_key_file.write(base64.b64decode(self.key_content)) + tmp_key_file.flush() + tmp_key_file.close() + os.chmod(tmp_key_file.name, 0o600) + 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 +92,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) diff --git a/tests/test_gcs_exporter.py b/tests/test_gcs_exporter.py new file mode 100644 index 0000000..eeb8fb4 --- /dev/null +++ b/tests/test_gcs_exporter.py @@ -0,0 +1,50 @@ +import base64 +import json +from unittest.mock import MagicMock, patch + +from murakami.exporters.gcs import GCSExporter + +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_key_content_uses_from_service_account_info(mock_client_cls): + """When key_content is in the config, credentials are loaded from inline + base64-encoded JSON rather than a file on disk.""" + mock_client_cls.from_service_account_info.return_value = MagicMock() + + GCSExporter( + name="test", + config={"target": "gs://bucket/path", "key_content": FAKE_KEY_CONTENT}, + ) + + 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_key_file_used_when_no_key_content(mock_client_cls): + """When key_content is absent, the exporter falls back to reading + credentials from the file path given by the key config option.""" + mock_client_cls.from_service_account_json.return_value = MagicMock() + + 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() diff --git a/tests/test_scp_exporter.py b/tests/test_scp_exporter.py new file mode 100644 index 0000000..7f875c4 --- /dev/null +++ b/tests/test_scp_exporter.py @@ -0,0 +1,67 @@ +import base64 +import os +from unittest.mock import MagicMock, patch + +from murakami.exporters.scp import SCPExporter + +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_key_content_writes_temp_file(mock_ssh_cls, mock_scp_cls): + """When key_content is in the config, the exporter decodes the base64 PEM, + writes it to a temporary file (mode 0o600), passes that path to + ssh.connect(), and deletes the file when done.""" + 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) + + exporter = SCPExporter( + name="test", + config={ + "target": "host:/remote/path", + "username": "user", + "key_content": FAKE_KEY_CONTENT, + }, + ) + 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, "key_filename should be set when key_content is provided" + 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_key_file_used_when_no_key_content(mock_ssh_cls, mock_scp_cls): + """When key_content is absent, the exporter uses the file path from the + key config option directly as key_filename.""" + 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) + + 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" diff --git a/utilities/encode_key.sh b/utilities/encode_key.sh new file mode 100755 index 0000000..72c4fb8 --- /dev/null +++ b/utilities/encode_key.sh @@ -0,0 +1,29 @@ +#!/bin/sh +# Encode a key file for use with Murakami's GCS or SCP exporter. +# +# Usage: +# encode_key.sh /path/to/sa.json (GCS service account JSON) +# encode_key.sh /path/to/id_rsa (SCP PEM private key) +# +# The encoded value is read by Murakami's existing environment variable +# config mechanism (load_env in __main__.py). Set it as: +# +# GCS exporter (in murakami.toml or as env var): +# [exporters.gcs] +# key_content = "" +# Or via environment variable: +# export MURAKAMI_EXPORTERS_GCS_KEY_CONTENT="" +# +# SCP exporter: +# [exporters.scp] +# key_content = "" +# Or via environment variable: +# export MURAKAMI_EXPORTERS_SCP_KEY_CONTENT="" + +if [ -z "$1" ]; then + echo "Usage: $0 /path/to/keyfile" >&2 + exit 1 +fi + +base64 -w 0 "$1" +echo ""