Merge branch 'master' into consistent-help

Author: bcwu

rsconnect/json_web_token.py

@@ -20,12 +20,31 @@
 BOOTSTRAP_SCOPE = "bootstrap"
 BOOTSTRAP_EXP = timedelta(minutes=15)
 
+SECRET_KEY_ENV = "CONNECT_BOOTSTRAP_SECRETKEY"
 
-def read_secret_key(keypath: str) -> bytes:
+
+def read_secret_key(keypath) -> bytes:
     """
     Reads a secret key as bytes given a path to a file containing a base64-encoded key.
+
+    The secret key can optionally be set with an environment variable.
     """
 
+    env_raw_data = os.getenv(SECRET_KEY_ENV)
+
+    if keypath is not None and env_raw_data is not None:
+        raise RSConnectException("Cannot specify secret key using both a keyfile and environment variable.")
+
+    if keypath is None and env_raw_data is None:
+        raise RSConnectException("Must specify secret key using either a keyfile or environment variable.")
+
+    # check if secret key was specified using an env variable first
+    if env_raw_data is not None:
+        try:
+            return base64.b64decode(env_raw_data.encode("utf-8"))
+        except binascii.Error:
+            raise RSConnectException("Unable to decode base64 data from environment variable: " + SECRET_KEY_ENV)
+
     if not os.path.exists(keypath):
         raise RSConnectException("Keypath does not exist.")
 

rsconnect/main.py

@@ -318,7 +318,6 @@ def _test_rstudio_creds(server: api.RStudioServer):
 @click.option(
     "--jwt-keypath",
     "-j",
-    required=True,
     help="The path to the file containing the private key used to sign the JWT.",
 )
 @click.option("--raw", "-r", is_flag=True, help="Return the API key as raw output rather than a JSON object")

tests/test_json_web_token.py

@@ -15,6 +15,7 @@
     BOOTSTRAP_SCOPE,
     DEFAULT_ISSUER,
     DEFAULT_AUDIENCE,
+    SECRET_KEY_ENV,
     read_secret_key,
     produce_bootstrap_output,
     is_jwt_compatible_python_version,
@@ -48,6 +49,8 @@ def setUp(self):
         # decoded copy of the base64-encoded key in testdata/jwt/secret.key
         self.secret_key = b"12345678901234567890123456789012345"
         self.secret_key_b64 = b"MTIzNDU2Nzg5MDEyMzQ1Njc4OTAxMjM0NTY3ODkwMTIzNDU="
+        # the environment variable version of the secret key will be stored as a string
+        self.secret_key_b64_env = "MTIzNDU2Nzg5MDEyMzQ1Njc4OTAxMjM0NTY3ODkwMTIzNDU="
 
     def assert_bootstrap_jwt_is_valid(self, payload, current_datetime):
         """
@@ -84,12 +87,38 @@ def test_read_secret_key(self):
         empty = read_secret_key("tests/testdata/jwt/empty_secret.key")
         self.assertEqual(empty, b"")
 
+        # might pass 'None' if we're using an environment variable
+        with pytest.raises(RSConnectException):
+            read_secret_key(None)
+
         with pytest.raises(RSConnectException):
             read_secret_key("invalid/path.key")
 
         with pytest.raises(RSConnectException):
             read_secret_key("tests/testdata/jwt/invalid_secret.key")
 
+        # environment variable replaces the need for a filepath
+        os.environ[SECRET_KEY_ENV] = self.secret_key_b64_env
+
+        valid_env = read_secret_key(None)
+        self.assertEqual(valid_env, self.secret_key)
+
+        # with env variable set, can't also attempt to read from file
+        with pytest.raises(RSConnectException):
+            read_secret_key("tests/testdata/jwt/secret.key")
+
+        with pytest.raises(RSConnectException):
+            read_secret_key("tests/testdata/jwt/empty_secret.key")
+
+        os.environ[SECRET_KEY_ENV] = "this_is_not_base64"
+
+        # env variable must also be a base64-encoded secret
+        with pytest.raises(RSConnectException):
+            read_secret_key(None)
+
+        # cleanup
+        del os.environ[SECRET_KEY_ENV]
+
     def test_validate_hs256_secret_key(self):
 
         with pytest.raises(RSConnectException):

tests/test_main.py

@@ -9,7 +9,7 @@
 import pytest
 from click.testing import CliRunner
 
-from rsconnect.json_web_token import is_jwt_compatible_python_version
+from rsconnect.json_web_token import SECRET_KEY_ENV, is_jwt_compatible_python_version
 
 from .utils import (
     apply_common_args,
@@ -567,6 +567,7 @@ def setUp(self):
         self.mock_server = "http://localhost:8080"
         self.mock_uri = "http://localhost:8080/__api__/v1/experimental/bootstrap"
         self.jwt_keypath = "tests/testdata/jwt/secret.key"
+        self.jwt_env_secret = "MTIzNDU2Nzg5MDEyMzQ1Njc4OTAxMjM0NTY3ODkwMTIzNDU="
 
         self.default_cli_args = [
             "bootstrap",
@@ -617,6 +618,39 @@ def test_bootstrap(self):
         expected_output = json.loads(open("tests/testdata/initial-admin-responses/success.json", "r").read())
         self.assertEqual(json_output, expected_output)
 
+    @httpretty.activate(verbose=True, allow_net_connect=False)
+    def test_bootstrap_env_var(self):
+        """
+        Normal initial-admin operation if secret key is configured using an environment variable
+        """
+        cli_args = [
+            "bootstrap",
+            "--server",
+            self.mock_server,
+            "--insecure",
+        ]
+
+        callback = self.create_bootstrap_mock_callback(200, {"api_key": "testapikey123"})
+
+        httpretty.register_uri(
+            httpretty.POST,
+            self.mock_uri,
+            body=callback,
+        )
+
+        os.environ[SECRET_KEY_ENV] = self.jwt_env_secret
+
+        runner = CliRunner()
+        result = runner.invoke(cli, cli_args)
+
+        self.assertEqual(result.exit_code, 0, result.output)
+
+        json_output = json.loads(result.output)
+        expected_output = json.loads(open("tests/testdata/initial-admin-responses/success.json", "r").read())
+        self.assertEqual(json_output, expected_output)
+
+        del os.environ[SECRET_KEY_ENV]
+
     @httpretty.activate(verbose=True, allow_net_connect=False)
     def test_bootstrap_misc_error(self):
         """
@@ -742,9 +776,46 @@ def test_bootstrap_invalid_server(self):
         )
 
     def test_boostrap_missing_jwt_option(self):
+        """
+        If jwt keyfile is not specified, it needs to be set using an environment variable
+        """
+        runner = CliRunner()
+        result = runner.invoke(cli, ["bootstrap", "--server", "http://a_server"])
+        self.assertEqual(result.exit_code, 1, result.output)
+        self.assertEqual(
+            result.output, "Error: Must specify secret key using either a keyfile or environment variable.\n"
+        )
+
+    def test_bootstrap_conflicting_jwt_option(self):
+        """
+        If jwt keyfile is specified, it cannot also be set using an environment variable
+        """
+
+        os.environ[SECRET_KEY_ENV] = "a_value"
+        runner = CliRunner()
+        result = runner.invoke(cli, self.default_cli_args)
+        self.assertEqual(result.exit_code, 1, result.output)
+        self.assertEqual(
+            result.output, "Error: Cannot specify secret key using both a keyfile and environment variable.\n"
+        )
+
+        del os.environ[SECRET_KEY_ENV]
+
+    def test_bootstrap_invalid_env_secret_key(self):
+        """
+        If jwt env variable is specified, it needs to be a valid base64-encoded value
+        """
+
+        os.environ[SECRET_KEY_ENV] = "a_value"
         runner = CliRunner()
         result = runner.invoke(cli, ["bootstrap", "--server", "http://a_server"])
-        self.assertEqual(result.exit_code, 2, result.output)
+        self.assertEqual(result.exit_code, 1, result.output)
+        self.assertEqual(
+            result.output,
+            "Error: Unable to decode base64 data from environment variable: CONNECT_BOOTSTRAP_SECRETKEY\n",
+        )
+
+        del os.environ[SECRET_KEY_ENV]
 
     @httpretty.activate(verbose=True, allow_net_connect=False)
     def test_bootstrap_raw_output(self):