diff --git a/kbcstorage/workspaces.py b/kbcstorage/workspaces.py index 4638d1b..83d4643 100644 --- a/kbcstorage/workspaces.py +++ b/kbcstorage/workspaces.py @@ -6,12 +6,45 @@ .. _here: http://docs.keboola.apiary.io/#reference/workspaces/ """ +from cryptography.hazmat.primitives import serialization +from cryptography.hazmat.primitives.asymmetric import rsa + from kbcstorage.base import Endpoint from kbcstorage.files import Files from kbcstorage.jobs import Jobs +from kbcstorage.tokens import Tokens from typing import List # the legacy Workspaces class below unfortunately defines its own method called list +BACKEND_SNOWFLAKE = 'snowflake' +LOGIN_TYPE_DEFAULT = 'default' +LOGIN_TYPE_SNOWFLAKE_SERVICE_KEYPAIR = 'snowflake-service-keypair' + +# sentinel distinguishing "not resolved yet" from "project has no default backend" +_DEFAULT_BACKEND_UNRESOLVED = object() + + +def _generate_rsa_key_pair(): + """ + Generate an RSA-2048 key pair for Snowflake key-pair authentication. + + Returns: + (private_key_pem, public_key_pem): Both keys PEM-encoded, the private + key in PKCS#8 format as expected by Snowflake drivers. + """ + private_key = rsa.generate_private_key(public_exponent=65537, key_size=2048) + private_key_pem = private_key.private_bytes( + encoding=serialization.Encoding.PEM, + format=serialization.PrivateFormat.PKCS8, + encryption_algorithm=serialization.NoEncryption(), + ).decode('ascii') + public_key_pem = private_key.public_key().public_bytes( + encoding=serialization.Encoding.PEM, + format=serialization.PublicFormat.SubjectPublicKeyInfo, + ).decode('ascii') + return private_key_pem, public_key_pem + + def _make_body(mapping, source_key='source', preserve: bool = True): """ Given a dict mapping Keboola tables to aliases, construct the body of @@ -44,6 +77,7 @@ def __init__(self, root_url, token): token (:obj:`str`): A storage API key. """ super().__init__(root_url, 'workspaces', token) + self._default_backend = _DEFAULT_BACKEND_UNRESOLVED def list(self): """ @@ -61,8 +95,9 @@ def detail(self, workspace_id): """ Retrieves information about a given workspace. - Note that the password to the workspace can only be retrieved when the - workspace is created. + Note that the workspace credentials (password or private key, + depending on the login type) are only available when the workspace + is created and cannot be retrieved later. Args: workspace_id (int or str): The id of the workspace. @@ -77,24 +112,78 @@ def create(self, backend=None, timeout=None, login_type=None, public_key=None, r """ Create a new Workspace and return the credentials. + On the snowflake backend, an omitted (or 'default') login_type would + create a deprecated password-based workspace. The client therefore + defaults to the 'snowflake-service-keypair' login type: when no + public_key is supplied, an RSA key pair is generated locally, the + public key is sent to the API and the private key is returned in + response['connection']['privateKey'] (it never leaves the client + otherwise and cannot be retrieved later). To get the deprecated + password-based workspace, pass login_type='snowflake-legacy-service' + explicitly. + Args: backend (:obj:`str`): The type of engine for the workspace. 'redshift', 'snowflake' or 'synapse'. Defaults to the project's default backend. timeout (int): The timeout, in seconds, for SQL statements. Only supported by snowflake backends. + login_type (:obj:`str`): The login type of the workspace, e.g. + 'snowflake-service-keypair', 'snowflake-person-keypair', + 'snowflake-legacy-service' or 'none'. Defaults to + 'snowflake-service-keypair' on snowflake, otherwise to the + backend's default. + public_key (:obj:`str`): PEM-encoded RSA public key to use with + key-pair login types. When omitted for the default snowflake + key-pair login, a key pair is generated locally. + read_all_objects (bool): Grant the workspace read-only access to + all project data. Raises: requests.HTTPError: If the API request fails. """ + private_key = None + effective_backend = backend or self._get_default_backend() + if effective_backend is None and login_type is not None: + raise ValueError( + "Cannot resolve the project's default backend from the token; " + "pass backend explicitly when using login_type." + ) + if effective_backend == BACKEND_SNOWFLAKE: + if login_type in (None, LOGIN_TYPE_DEFAULT): + login_type = LOGIN_TYPE_SNOWFLAKE_SERVICE_KEYPAIR + if login_type == LOGIN_TYPE_SNOWFLAKE_SERVICE_KEYPAIR and public_key is None: + private_key, public_key = _generate_rsa_key_pair() + if login_type is not None: + # the API rejects loginType without an explicit backend + backend = effective_backend + body = { - 'backend': backend, - 'statementTimeoutSeconds': timeout, - 'loginType': login_type, - 'publicKey': public_key, - 'readOnlyStorageAccess': str(read_all_objects).lower() # convert bool to lowercase true or false + k: v for k, v in { + 'backend': backend, + 'statementTimeoutSeconds': timeout, + 'loginType': login_type, + 'publicKey': public_key, + 'readOnlyStorageAccess': str(read_all_objects).lower() # convert bool to lowercase true or false + }.items() + if v is not None } - return self._post(self.base_url, data=body) + response = self._post(self.base_url, data=body) + if private_key is not None: + response.setdefault('connection', {})['privateKey'] = private_key + return response + + def _get_default_backend(self): + """ + Resolve the project's default backend from the token detail. + + The value is cached on the instance - a project's default backend is + effectively immutable for the client's lifetime. + """ + if self._default_backend is _DEFAULT_BACKEND_UNRESOLVED: + token_info = Tokens(self.root_url, self.token).verify() + self._default_backend = (token_info.get('owner') or {}).get('defaultBackend') + return self._default_backend def delete(self, workspace_id): """ @@ -116,6 +205,11 @@ def reset_password(self, workspace_id): """ Generate a new password for the workspace. + Only supported for password-based login types (e.g. the deprecated + 'snowflake-legacy-service'). For key-pair workspaces rotate the + credentials with set_public_key() using a freshly generated key pair + instead. + Args: workspace_id (int or str): The id of the workspace for which the password should be reset. diff --git a/pyproject.toml b/pyproject.toml index a6f05e2..0d5b328 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -21,6 +21,7 @@ classifiers = [ dependencies = [ "boto3", "azure-storage-blob", + "cryptography", "urllib3<2.0.0", # Frozen until fixed: https://github.com/boto/botocore/issues/2926 # Dev dependencies "requests", diff --git a/tests/mocks/test_workspaces.py b/tests/mocks/test_workspaces.py index b343d46..f598562 100644 --- a/tests/mocks/test_workspaces.py +++ b/tests/mocks/test_workspaces.py @@ -1,14 +1,19 @@ """ -Asses basic functionality of the Workspace endpoint. +Assess basic functionality of the Workspace endpoint. """ +import copy import unittest +from urllib.parse import parse_qs + import responses from requests import HTTPError from kbcstorage.workspaces import Workspaces +from .token_responses import verify_token_response from .workspace_responses import (list_response, detail_response, load_tables_response, create_response, + keypair_create_response, reset_password_response) @@ -91,18 +96,191 @@ def test_detail_inexsitent_workspace(self): @responses.activate def test_create(self): """ - Workspace endpoint mock creates new workspace + Workspace endpoint mock creates new workspace. With no backend given, + the project default backend (snowflake) is resolved from the token and + a key-pair workspace is created instead of a password one. """ + responses.add( + responses.Response( + method='GET', + url='https://connection.keboola.com/v2/storage/tokens/verify', + json=verify_token_response + ) + ) responses.add( responses.Response( method='POST', url='https://connection.keboola.com/v2/storage/workspaces', - json=create_response + json=keypair_create_response ) ) created_detail = self.ws.create() + request_body = parse_qs(responses.calls[1].request.body, keep_blank_values=True) + assert request_body['backend'] == ['snowflake'] + assert request_body['loginType'] == ['snowflake-service-keypair'] + assert 'BEGIN PUBLIC KEY' in request_body['publicKey'][0] + assert 'BEGIN PRIVATE KEY' in created_detail['connection']['privateKey'] + + @responses.activate + def test_create_snowflake_defaults_to_keypair(self): + """ + With an explicit snowflake backend no token verify is needed and the + login type defaults to snowflake-service-keypair with a locally + generated key pair. + """ + responses.add( + responses.Response( + method='POST', + url='https://connection.keboola.com/v2/storage/workspaces', + json=keypair_create_response + ) + ) + created_detail = self.ws.create(backend='snowflake') + assert len(responses.calls) == 1 + request_body = parse_qs(responses.calls[0].request.body, keep_blank_values=True) + assert request_body['loginType'] == ['snowflake-service-keypair'] + assert 'BEGIN PUBLIC KEY' in request_body['publicKey'][0] + assert 'BEGIN PRIVATE KEY' in created_detail['connection']['privateKey'] + + @responses.activate + def test_create_snowflake_with_own_public_key(self): + """ + A caller-supplied public key is passed through and no private key is + injected into the response. + """ + responses.add( + responses.Response( + method='POST', + url='https://connection.keboola.com/v2/storage/workspaces', + json=keypair_create_response + ) + ) + created_detail = self.ws.create(backend='snowflake', public_key='my-public-key') + request_body = parse_qs(responses.calls[0].request.body, keep_blank_values=True) + assert request_body['loginType'] == ['snowflake-service-keypair'] + assert request_body['publicKey'] == ['my-public-key'] + assert 'privateKey' not in created_detail['connection'] + + @responses.activate + def test_create_snowflake_explicit_legacy_login_type(self): + """ + An explicitly requested password login type is passed through + unchanged and no key pair is generated. + """ + responses.add( + responses.Response( + method='POST', + url='https://connection.keboola.com/v2/storage/workspaces', + json=create_response + ) + ) + created_detail = self.ws.create(backend='snowflake', login_type='snowflake-legacy-service') + request_body = parse_qs(responses.calls[0].request.body, keep_blank_values=True) + assert request_body['loginType'] == ['snowflake-legacy-service'] + assert 'publicKey' not in request_body assert created_detail['connection']['password'] == 'abc' + @responses.activate + def test_create_non_snowflake_default_backend_unchanged(self): + """ + When the project default backend resolved from the token is not + snowflake, no login type, public key or backend is sent. + """ + bigquery_verify_response = copy.deepcopy(verify_token_response) + bigquery_verify_response['owner']['defaultBackend'] = 'bigquery' + responses.add( + responses.Response( + method='GET', + url='https://connection.keboola.com/v2/storage/tokens/verify', + json=bigquery_verify_response + ) + ) + responses.add( + responses.Response( + method='POST', + url='https://connection.keboola.com/v2/storage/workspaces', + json=create_response + ) + ) + self.ws.create() + request_body = parse_qs(responses.calls[1].request.body, keep_blank_values=True) + assert 'backend' not in request_body + assert 'loginType' not in request_body + assert 'publicKey' not in request_body + + @responses.activate + def test_create_caches_default_backend(self): + """ + The default backend is resolved via the token verify call only once + per endpoint instance. + """ + responses.add( + responses.Response( + method='GET', + url='https://connection.keboola.com/v2/storage/tokens/verify', + json=verify_token_response + ) + ) + responses.add( + responses.Response( + method='POST', + url='https://connection.keboola.com/v2/storage/workspaces', + json=keypair_create_response + ) + ) + responses.add( + responses.Response( + method='POST', + url='https://connection.keboola.com/v2/storage/workspaces', + json=keypair_create_response + ) + ) + self.ws.create() + self.ws.create() + verify_calls = [c for c in responses.calls if c.request.method == 'GET'] + assert len(verify_calls) == 1 + + @responses.activate + def test_create_login_type_without_resolvable_backend_raises(self): + """ + The API rejects loginType without an explicit backend, so the client + raises a clear error when the default backend cannot be resolved. The + missing value is cached too - verify is not re-queried on retry. + """ + no_backend_verify_response = copy.deepcopy(verify_token_response) + del no_backend_verify_response['owner']['defaultBackend'] + responses.add( + responses.Response( + method='GET', + url='https://connection.keboola.com/v2/storage/tokens/verify', + json=no_backend_verify_response + ) + ) + with self.assertRaises(ValueError): + self.ws.create(login_type='none') + with self.assertRaises(ValueError): + self.ws.create(login_type='none') + assert len(responses.calls) == 1 + + @responses.activate + def test_create_non_snowflake_backend_unchanged(self): + """ + Non-snowflake backends keep the original behavior - no login type or + public key is sent. + """ + responses.add( + responses.Response( + method='POST', + url='https://connection.keboola.com/v2/storage/workspaces', + json=create_response + ) + ) + self.ws.create(backend='bigquery') + request_body = parse_qs(responses.calls[0].request.body, keep_blank_values=True) + assert request_body['backend'] == ['bigquery'] + assert 'loginType' not in request_body + assert 'publicKey' not in request_body + @responses.activate def test_delete(self): """ diff --git a/tests/mocks/workspace_responses.py b/tests/mocks/workspace_responses.py index cfebf50..ac5c6bc 100644 --- a/tests/mocks/workspace_responses.py +++ b/tests/mocks/workspace_responses.py @@ -75,6 +75,14 @@ } } +keypair_create_response = { + **create_response, + "connection": { + **{k: v for k, v in create_response["connection"].items() if k != "password"}, + "loginType": "snowflake-service-keypair", + }, +} + load_tables_response = { "id": 22077337, "status": "waiting",