diff --git a/snowflake/assets/configuration/spec.yaml b/snowflake/assets/configuration/spec.yaml index fbe56db15f604..1bf7bd2a1183e 100644 --- a/snowflake/assets/configuration/spec.yaml +++ b/snowflake/assets/configuration/spec.yaml @@ -114,6 +114,20 @@ files: value: type: string example: /path/to/token + - name: private_key_path + description: | + The path to the file that contains the private key used to connect to Snowflake. + The key is re-read at every check run. + value: + type: string + display_default: null + example: /path/to/private_key + - name: private_key_password + secret: true + description: | + The password protecting the private key file in the `private_key_path` option. + value: + type: string - name: client_session_keep_alive description: | If set to true, Snowflake keeps the session active indefinitely as long as the connection is active, diff --git a/snowflake/datadog_checks/snowflake/check.py b/snowflake/datadog_checks/snowflake/check.py index 93bc6473fc0fe..483029c076c23 100644 --- a/snowflake/datadog_checks/snowflake/check.py +++ b/snowflake/datadog_checks/snowflake/check.py @@ -4,6 +4,8 @@ from contextlib import closing import snowflake.connector as sf +from cryptography.hazmat.backends import default_backend +from cryptography.hazmat.primitives import serialization from datadog_checks.base import AgentCheck, ConfigurationError, to_native_string from datadog_checks.base.utils.db import QueryManager @@ -77,10 +79,30 @@ def __init__(self, *args, **kwargs): self._query_manager = QueryManager(self, self.execute_query_raw, queries=self.metric_queries, tags=self._tags) self.check_initializations.append(self._query_manager.compile_queries) - def renew_token(self): - self.log.debug("Renewing Snowflake client token") - with open(self._config.token_path, 'rb', encoding="UTF-8") as f: - self._config.token = f.read() + def read_token(self): + if self._config.token_path: + self.log.debug("Renewing Snowflake client token") + with open(self._config.token_path, 'rb', encoding="UTF-8") as f: + self._config.token = f.read() + + def read_key(self): + if self._config.private_key_path: + self.log.debug("Reading Snowflake client key for key pair authentication") + # https://docs.snowflake.com/en/user-guide/python-connector-example.html#using-key-pair-authentication-key-pair-rotation + with open(self._config.private_key_path, "rb") as key: + p_key = serialization.load_pem_private_key( + key.read(), password=self._config.private_key_password, backend=default_backend() + ) + + pkb = p_key.private_bytes( + encoding=serialization.Encoding.DER, + format=serialization.PrivateFormat.PKCS8, + encryption_algorithm=serialization.NoEncryption(), + ) + + return pkb + + return None def check(self, _): if self.instance.get('user'): @@ -126,8 +148,7 @@ def connect(self): self.proxy_port, ) - if self._config.token_path: - self.renew_token() + self.read_token() try: conn = sf.connect( @@ -145,6 +166,7 @@ def connect(self): ocsp_response_cache_filename=self._config.ocsp_response_cache_filename, authenticator=self._config.authenticator, token=self._config.token, + private_key=self.read_key(), client_session_keep_alive=self._config.client_keep_alive, proxy_host=self.proxy_host, proxy_port=self.proxy_port, diff --git a/snowflake/datadog_checks/snowflake/config.py b/snowflake/datadog_checks/snowflake/config.py index e180032921ee2..a2b4bc575de85 100644 --- a/snowflake/datadog_checks/snowflake/config.py +++ b/snowflake/datadog_checks/snowflake/config.py @@ -3,7 +3,7 @@ # Licensed under a 3-clause BSD style license (see LICENSE) from typing import List, Optional -from datadog_checks.base import ConfigurationError, is_affirmative +from datadog_checks.base import ConfigurationError, ensure_bytes, is_affirmative class Config(object): @@ -19,7 +19,7 @@ class Config(object): 'snowflake.logins', ] - AUTHENTICATION_MODES = ['snowflake', 'oauth'] + AUTHENTICATION_MODES = ['snowflake', 'oauth', 'snowflake_jwt'] def __init__(self, instance=None): if instance is None: @@ -41,6 +41,8 @@ def __init__(self, instance=None): authenticator = instance.get('authenticator', 'snowflake') token = instance.get('token', None) token_path = instance.get('token_path', None) + private_key_path = instance.get('private_key_path', None) + private_key_password = ensure_bytes(instance.get('private_key_password', None)) client_keep_alive = instance.get('client_session_keep_alive', False) aggregate_last_24_hours = instance.get('aggregate_last_24_hours', False) @@ -56,8 +58,8 @@ def __init__(self, instance=None): if is_affirmative(passcode_in_password) and passcode is None: raise ConfigurationError('MFA enabled, please specify a passcode') - if password is None: - raise ConfigurationError('Must specify a password if using snowflake authentication') + if password is None and private_key_path is None: + raise ConfigurationError('Must specify a password or a private key if using snowflake authentication') elif authenticator == 'oauth': if token is None and token_path is None: @@ -92,5 +94,7 @@ def __init__(self, instance=None): self.authenticator = authenticator self.token = token self.token_path = token_path + self.private_key_path = private_key_path + self.private_key_password = private_key_password self.client_keep_alive = client_keep_alive self.aggregate_last_24_hours = aggregate_last_24_hours diff --git a/snowflake/datadog_checks/snowflake/config_models/defaults.py b/snowflake/datadog_checks/snowflake/config_models/defaults.py index 89ccae578127a..0978291fe8d42 100644 --- a/snowflake/datadog_checks/snowflake/config_models/defaults.py +++ b/snowflake/datadog_checks/snowflake/config_models/defaults.py @@ -90,6 +90,14 @@ def instance_password(field, value): return get_default_field_value(field, value) +def instance_private_key_password(field, value): + return get_default_field_value(field, value) + + +def instance_private_key_path(field, value): + return get_default_field_value(field, value) + + def instance_schema_(field, value): return 'ACCOUNT_USAGE' diff --git a/snowflake/datadog_checks/snowflake/config_models/instance.py b/snowflake/datadog_checks/snowflake/config_models/instance.py index 6fa5cf0547b7c..6fb2a0e8c0677 100644 --- a/snowflake/datadog_checks/snowflake/config_models/instance.py +++ b/snowflake/datadog_checks/snowflake/config_models/instance.py @@ -47,6 +47,8 @@ class Config: ocsp_response_cache_filename: Optional[str] only_custom_queries: Optional[bool] password: Optional[str] + private_key_password: Optional[str] + private_key_path: Optional[str] role: str schema_: Optional[str] = Field(None, alias='schema') service: Optional[str] diff --git a/snowflake/datadog_checks/snowflake/config_models/validators.py b/snowflake/datadog_checks/snowflake/config_models/validators.py index 96996e5bfa2e2..9a6d2b7d5a72b 100644 --- a/snowflake/datadog_checks/snowflake/config_models/validators.py +++ b/snowflake/datadog_checks/snowflake/config_models/validators.py @@ -8,4 +8,10 @@ def initialize_instance(values, **kwargs): if 'username' not in values and 'user' in values: values['username'] = values['user'] + if 'private_key_password' in values and 'private_key_path' not in values: + raise Exception( + 'Option `private_key_password` is set but not option `private_key_path`. ' + 'Set `private_key_path` or remove `private_key_password` entry.' + ) + return values diff --git a/snowflake/datadog_checks/snowflake/data/conf.yaml.example b/snowflake/datadog_checks/snowflake/data/conf.yaml.example index b75f035990424..e00b122868114 100644 --- a/snowflake/datadog_checks/snowflake/data/conf.yaml.example +++ b/snowflake/datadog_checks/snowflake/data/conf.yaml.example @@ -122,6 +122,17 @@ instances: # # token_path: /path/to/token + ## @param private_key_path - string - optional + ## The path to the file that contains the private key used to connect to Snowflake. + ## The key is re-read at every check run. + # + # private_key_path: /path/to/private_key + + ## @param private_key_password - string - optional + ## The password protecting the private key file in the `private_key_path` option. + # + # private_key_password: + ## @param client_session_keep_alive - boolean - optional - default: false ## If set to true, Snowflake keeps the session active indefinitely as long as the connection is active, ## even if there is no activity from the user. diff --git a/snowflake/tests/common.py b/snowflake/tests/common.py index 12b44946e6383..c1c4e77eb1076 100644 --- a/snowflake/tests/common.py +++ b/snowflake/tests/common.py @@ -42,5 +42,8 @@ 'pip install --upgrade setuptools', 'pip install /home/snowflake_connector_patch', ], - 'docker_volumes': ['{}:/home/snowflake_connector_patch'.format(os.path.join(HERE, 'snowflake_connector_patch'))], + 'docker_volumes': [ + '{}:/home/snowflake_connector_patch'.format(os.path.join(HERE, 'snowflake_connector_patch')), + '{}:/home/keys'.format(os.path.join(HERE, 'keys')), + ], } diff --git a/snowflake/tests/keys/rsa_key_example.p8 b/snowflake/tests/keys/rsa_key_example.p8 new file mode 100644 index 0000000000000..6d2439379a00a --- /dev/null +++ b/snowflake/tests/keys/rsa_key_example.p8 @@ -0,0 +1,28 @@ +-----BEGIN PRIVATE KEY----- +MIIEvAIBADANBgkqhkiG9w0BAQEFAASCBKYwggSiAgEAAoIBAQCXeqIH2MXr6E7p +utpe+ZC/31oMMgXi7EPqjeNdbEMkM5c1hCLSyy762mlafkvxqEaL+OIEamZXDtOq +/cenVogdu5gvP/mMkijxWnMg7CNOtg6ulME2UqOPiMRQRJm9iXFoqx1xLSVSUh08 +0ydbh2k8+yav/eRqaAJyREOE0IvafyeG1ocHTL/D6eHos3epB9T4OLXS6UjOvLoJ +9eu4GLj7R6GuMB1X15RlmYP+XLWWO2pKjlfBdY8KxQ48Q5GQHpKV10ytuH36VO0q +UvGRy5svzQEcI7VX/HCLkiFtbL7yQ+gyKACYql2CpBaBR0IDYAyx3KjrFdLSVkZ+ +goWNcTThAgMBAAECggEAYPKDsTHzukA8ASsth4uCMMrp/tQlIE7GSN/2nFwlvI2o +QQAIqZCQyMvwkZIzWL6yJ/Np9BmE1kMPBWjW1ypyg1UE2XjAZk9FFPMmq/N1xXKP +gvyBjBrnw31s51KAcuX8R1j8xup18HHNdJhPoCzSFc1HvWtpPRDEQA2giOhQGc40 +xcz34161dj8BAeeESHaIfed7WTwBu1IBTpVmAKPJudlAf3EjiAcO/pclLS7tJwN5 +qsX0K3oHkm4W1aaOLo/G4akUxJ72bS8ny9w8wvJmJKIJ93koHwJ5HnThL5UjGhCD +cAJPCdUYm1oUtKAsTkIh2SC/sGBH28Q20D5UGHms4QKBgQD4UgvXIQYIcUv0M0qA +vGxyRDsG474Q8PqnClRAEZc/NA3S0g/B+jq0kyq7J3DLBtLCDEPmjBKfAbm39D2t +TE4sebo2/CTYmJy5FgiJxdSaEoZOrPAZAJfpkN3lfohN6McIT+0YUTvdgnDk/5xo +NwPQQEfdN9OtWMqvaFMrr6gf8wKBgQCcKeSH2/t39r79LmZvxW42+dBQvhoDVR4J +OsTYx058YWdAbWBy2CccCkDdOQnuPBsqIkQy5YQ4E6dt5AqOP9Kcl4xrfyocvQpy +ULa/sLl5xZTtKf6iyRhOuZS+LSgychoc3Y50xyNKoVAMv2ddvQnEQrEDQizbO6gW +43idKEKg2wKBgAzukOlKMfs8kz0LcsTTiz5EKWLJd3uAYT1Tv2F6yQqklle1UtbC +Rk5jH6WRf0FDgLRUWTDneIzJVTesQ44D3EpaqIT2iqCxCfBloloycEj5z/7G6NYU +ftTOE5BBD64nAj5/kxRiHqEBiwmR+j4/JzawMk3l+2MarauG3lX3FuVbAoGADoVe +uLtd4MPS8pvz7oS/QOFt23Qx2wl5J4aNc1LlG2+7OCRziXpL+LGDYo7BO6PfKsXQ +7aKl7sj1EqTXzm5k2SbGaeCDO/TgGc0jkSOPu6EBviPfh6eHWRqsmBp+2GH/x5ta +ecVipLfnR6gspmzDkbpZ12G55hDgCnDQcFykBW0CgYBeYRRfvZuJVDbA7M3DygfI +DCNPpgcjaECt5Uz/3N0udBNsoXY9vRakRlhKANLYesMUd76mUjVWu6Kc+EJtyEcC +q2plTc4gowD/hhxT1Fw95p6yxz8m+eHFc8BWAaHtoGMMgK1/wm1UXYe/hQDu39g9 +ojckZpT64MHda9CvnF814g== +-----END PRIVATE KEY----- diff --git a/snowflake/tests/keys/rsa_key_pass_example.p8 b/snowflake/tests/keys/rsa_key_pass_example.p8 new file mode 100644 index 0000000000000..b176b5efc04e2 --- /dev/null +++ b/snowflake/tests/keys/rsa_key_pass_example.p8 @@ -0,0 +1,29 @@ +-----BEGIN ENCRYPTED PRIVATE KEY----- +MIIE6TAbBgkqhkiG9w0BBQMwDgQIH+eVVXditq8CAggABIIEyLIREkj8xviusabN +Eb9xkeuJJ2AMUHZjDlrhlUDN8LwZc5REcxvv1Y+aEvAaiVFjVM8NdhLLmreW6/mf +1P2Ms0p8dUXg4w0ulU0dfoMwDSTz4vy/F66+pRNhctcMCxsbUITtYwe5tFKffJuC +ovtuj0WbqgDngMiB+LUI9MDxBCKsS0tYoIE9RlNCplWZmNW8dY0rTAwJx0K0D5Hm +cxZkj4E95wn2Jw57jdnbw1tEZyDAmR+v1aKVCtD39V8rAvBTBlfPAtD5IPLtOnem +gvPv6EuSjwMmuKREBSkrZr/ibPbPcJwSvHBoQH1YVk5qQjXp+qXjrPA0xGFmYZ5O +fc4kbMhZfB9QbioVql3wBd1we+3au+lC+Xao0RniOaX3waVOvvlsZFRBCnULuAVb +YdezqDTpukxhi1Zp9+nXNFf0jQdKbEAcMP9MMe14IGTH0Zqs9g63jzOig5V1wU3a +/ghalyC9RxnICimo/JObRl3hIgLK2VSoLIxguRdrScKXAFv2RCk6QuoavJwoxfdc +ElXhHYuXM6+QzCHfaT3G1NPg9aJ40Ycp08Js/5DLzVjhdwBtQuJCpO+c70aIM5qw +zb/4ZvOU+vFQVZV5cEYJEuHBvQ65wnJPhrt6tWi87xlW2Enbx9jNKThMeT6UtV6M +Rapy8nQp7Tg48PkuK6m7FrnBOrf6g+QmqpEuO+G7wH4b+CiLzcGiivmg3LkNXBO3 +YRzY1xYLemowwdtLKFVMw4otXQqjRey6TU54MNLMGLg8kmqBmOi0ti+noInWMsT7 +sohpq7KhZtsDICl2sTxSz8UKZJG8vIuhBGoVfl3YKtp+M5yTB4/uGi3urV4jg266 +FRRfvEvRFh6Bvu016L0ClXXnsOyqcd9vAie+bQUY39L/WTetMh8U/9AKDUUvJONa +NJf4wnnlhfyVx8Tg2DoPf7V0ekK2fPfmvxhlM+RVPKw1GxXzum2q4TKPBZaZJAR6 +aGJwb5FtfM6EmTa5vJBGOxfOf9cVAu67OoMAyAXYQDFM+ruOXzTrIa5No67FRXkt +LUWgRClSYPg0tEeBiNut4M5ZyxIwNBKG3Y3gFyUbMVpj8EfnkmAYTUJCKeXKnor6 +JmtAAYYGfugBzdyZ74qYdEqai1cMkTE2NNm2xeJBrq5Cg3e1qWnnJk8+OZkFAxNF +i2eNEzcgwQDhz40M/3gPzEjYtQmaVibtN/kN4/UMGYHrKUIqZ4HeddouZf30wXo2 +ovjhFiXJBE2HgDYDJ42rCDhgMWir9XFgsM8+51Xg0sNBDIAa2PRhnUM46jLMWZo/ +PU/6vK1jKPejeuXAeZwGSFJOdlC7U7X4+zZ8gBtD7YiiqJMmKl/0Go6ACXokB7r7 +i097/NOpH0Q0z4TsC8xEHinDTepmTskAaEHEFWrULQ3kd5I3sym2OHRz/wrxOz6v +9IWshaEkgp+UmdPeh0qiZwP6D9wuRSBeb5RFTfPebVlDnUbPvcItaI8oS1ux6WJe +SOA16NvCjeyjjnWKFKKYV2VFbIcZevO+GwsPh0zmtPdAgb2dAATfjb1PLbV2yaCs +nAjIhDEsG54pmJA+MCdY/auJeDqu8OrdftLR6gCDqNRJOM7c0bM09dstkLXsiadm +Izh71I5VdjD088Wdeg== +-----END ENCRYPTED PRIVATE KEY----- diff --git a/snowflake/tests/test_snowflake.py b/snowflake/tests/test_snowflake.py index 4890261880f0d..13683e5ccc5cc 100644 --- a/snowflake/tests/test_snowflake.py +++ b/snowflake/tests/test_snowflake.py @@ -183,6 +183,7 @@ def test_token_path(dd_run_check, aggregator): 'ocsp_response_cache_filename': None, 'authenticator': 'oauth', 'client_session_keep_alive': False, + 'private_key': None, 'proxy_host': None, 'proxy_port': None, 'proxy_user': None, diff --git a/snowflake/tests/test_unit.py b/snowflake/tests/test_unit.py index 3787414b40d1e..3acf1a0e6b27a 100644 --- a/snowflake/tests/test_unit.py +++ b/snowflake/tests/test_unit.py @@ -2,6 +2,7 @@ # All rights reserved # Licensed under a 3-clause BSD style license (see LICENSE) import copy +import os import mock import pytest @@ -81,6 +82,7 @@ def test_default_auth(instance): ocsp_response_cache_filename=None, authenticator='snowflake', token=None, + private_key=None, client_session_keep_alive=False, proxy_host=None, proxy_port=None, @@ -113,6 +115,7 @@ def test_oauth_auth(oauth_instance): ocsp_response_cache_filename=None, authenticator='oauth', token='testtoken', + private_key=None, client_session_keep_alive=False, proxy_host=None, proxy_port=None, @@ -121,6 +124,54 @@ def test_oauth_auth(oauth_instance): ) +def test_key_auth(dd_run_check, instance): + # Key auth + inst = copy.deepcopy(instance) + inst['private_key_path'] = os.path.join(os.path.dirname(__file__), 'keys', 'rsa_key_example.p8') + + check = SnowflakeCheck(CHECK_NAME, {}, [inst]) + # Checking size instead of the read key + read_key = check.read_key() + assert len(read_key) == 1216 + + with mock.patch('datadog_checks.snowflake.check.sf') as sf: + check = SnowflakeCheck(CHECK_NAME, {}, [inst]) + dd_run_check(check) + sf.connect.assert_called_with( + user='testuser', + password='pass', + account='test_acct.us-central1.gcp', + database='SNOWFLAKE', + schema='ACCOUNT_USAGE', + warehouse=None, + role='ACCOUNTADMIN', + passcode_in_password=False, + passcode=None, + client_prefetch_threads=4, + login_timeout=3, + ocsp_response_cache_filename=None, + authenticator='snowflake', + token=None, + private_key=read_key, + client_session_keep_alive=False, + proxy_host=None, + proxy_port=None, + proxy_user=None, + proxy_password=None, + ) + + inst['private_key_path'] = os.path.join(os.path.dirname(__file__), 'keys', 'wrong_key.p8') + check = SnowflakeCheck(CHECK_NAME, {}, [inst]) + with pytest.raises(FileNotFoundError): + check.read_key() + + # Read key protected by a passphrase + inst['private_key_path'] = os.path.join(os.path.dirname(__file__), 'keys', 'rsa_key_pass_example.p8') + inst['private_key_password'] = 'keypass' + check = SnowflakeCheck(CHECK_NAME, {}, [inst]) + assert len(check.read_key()) == 1218 + + def test_proxy_settings(instance): init_config = { 'proxy_host': 'testhost', @@ -149,6 +200,7 @@ def test_proxy_settings(instance): ocsp_response_cache_filename=None, authenticator='snowflake', token=None, + private_key=None, client_session_keep_alive=False, proxy_host='testhost', proxy_port=8000,