diff --git a/azure-kusto-data/azure/kusto/data/client_base.py b/azure-kusto-data/azure/kusto/data/client_base.py index 418ecba1..6bc25239 100644 --- a/azure-kusto-data/azure/kusto/data/client_base.py +++ b/azure-kusto-data/azure/kusto/data/client_base.py @@ -9,7 +9,6 @@ from requests import Response, Session from azure.kusto.data._cloud_settings import CloudSettings -from azure.kusto.data._token_providers import CloudInfoTokenProvider from .client_details import ClientDetails from .client_request_properties import ClientRequestProperties from .exceptions import KustoServiceError, KustoThrottlingError, KustoApiError @@ -78,16 +77,14 @@ def set_proxy(self, proxy_url: str): def validate_endpoint(self): if not self._endpoint_validated and self._aad_helper is not None: - if isinstance(self._aad_helper.token_provider, CloudInfoTokenProvider): - endpoint = CloudSettings.get_cloud_info_for_cluster( + well_known_kusto_endpoints.validate_trusted_endpoint( + self._kusto_cluster, + lambda: CloudSettings.get_cloud_info_for_cluster( self._kusto_cluster, self._aad_helper.token_provider._proxy_dict, self._session if isinstance(self._session, Session) else None, - ).login_endpoint - well_known_kusto_endpoints.validate_trusted_endpoint( - self._kusto_cluster, - endpoint, - ) + ).login_endpoint, + ) self._endpoint_validated = True @staticmethod diff --git a/azure-kusto-data/azure/kusto/data/kusto_trusted_endpoints.py b/azure-kusto-data/azure/kusto/data/kusto_trusted_endpoints.py index dd5e4bd2..ca8bd31b 100644 --- a/azure-kusto-data/azure/kusto/data/kusto_trusted_endpoints.py +++ b/azure-kusto-data/azure/kusto/data/kusto_trusted_endpoints.py @@ -1,5 +1,5 @@ import copy -from typing import List, Dict +from typing import Callable, List, Dict, Union from urllib.parse import urlparse from azure.kusto.data.helpers import get_string_tail_lower_case @@ -73,23 +73,29 @@ def add_trusted_hosts(self, rules, replace): self._additional_matcher = create_fast_suffix_matcher_from_existing(rules, None if replace else self._additional_matcher) - def validate_trusted_endpoint(self, endpoint: str, login_endpoint: str): + def validate_trusted_endpoint(self, endpoint: str, login_endpoint: Union[str, Callable[[], str]]): hostname = urlparse(endpoint).hostname self.validate_hostname_is_trusted(hostname if hostname is not None else endpoint, login_endpoint) - def validate_hostname_is_trusted(self, hostname: str, login_endpoint: str): + def validate_hostname_is_trusted(self, hostname: str, login_endpoint: Union[str, Callable[[], str]]): if _is_local_address(hostname): return if self._override_matcher is not None: if self._override_matcher(hostname): return else: - matcher = self._matchers.get(login_endpoint.lower()) - if matcher is not None and matcher.is_match(hostname): + additional_matcher = self._additional_matcher + if additional_matcher is not None and additional_matcher.is_match(hostname): return - matcher = self._additional_matcher - if matcher is not None and matcher.is_match(hostname): + if any(matcher.is_match(hostname) for matcher in self._matchers.values()): + resolved_login_endpoint = login_endpoint() if callable(login_endpoint) else login_endpoint + matcher = self._matchers.get(resolved_login_endpoint.lower()) + if matcher is not None and matcher.is_match(hostname): + return + + additional_matcher = self._additional_matcher + if additional_matcher is not None and additional_matcher.is_match(hostname): return raise KustoClientInvalidConnectionStringException( diff --git a/azure-kusto-data/tests/test_endpoint_validation.py b/azure-kusto-data/tests/test_endpoint_validation.py new file mode 100644 index 00000000..db7a6029 --- /dev/null +++ b/azure-kusto-data/tests/test_endpoint_validation.py @@ -0,0 +1,115 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License +from unittest.mock import AsyncMock, Mock, patch + +import pytest + +from azure.kusto.data import KustoClient, KustoConnectionStringBuilder +from azure.kusto.data._cloud_settings import DEFAULT_PUBLIC_LOGIN_URL +from azure.kusto.data.aio import KustoClient as AsyncKustoClient +from azure.kusto.data.exceptions import KustoClientInvalidConnectionStringException +from azure.kusto.data.kusto_trusted_endpoints import KustoTrustedEndpoints, MatchRule + + +UNTRUSTED_HOST = "https://kusto.attacker.example.com" +TRUSTED_HOST = "https://somecluster.kusto.windows.net" +TOKEN = "a token that must not be transmitted" + + +def _sync_kcsb(authentication_method, token_callback): + if authentication_method == "user_token": + return KustoConnectionStringBuilder.with_aad_user_token_authentication(UNTRUSTED_HOST, TOKEN) + if authentication_method == "application_token": + return KustoConnectionStringBuilder.with_aad_application_token_authentication(UNTRUSTED_HOST, TOKEN) + return KustoConnectionStringBuilder.with_token_provider(UNTRUSTED_HOST, token_callback) + + +def _async_kcsb(authentication_method, token_callback): + if authentication_method == "user_token": + return KustoConnectionStringBuilder.with_aad_user_token_authentication(UNTRUSTED_HOST, TOKEN) + if authentication_method == "application_token": + return KustoConnectionStringBuilder.with_aad_application_token_authentication(UNTRUSTED_HOST, TOKEN) + return KustoConnectionStringBuilder.with_async_token_provider(UNTRUSTED_HOST, token_callback) + + +@pytest.mark.parametrize("authentication_method", ["user_token", "application_token", "callback"]) +def test_sync_authentication_rejects_untrusted_host_before_token_use(authentication_method): + token_callback = Mock(return_value=TOKEN) + kcsb = _sync_kcsb(authentication_method, token_callback) + + with ( + patch("azure.kusto.data.client_base.CloudSettings.get_cloud_info_for_cluster") as cloud_info, + patch("azure.kusto.data.security._get_header_from_dict") as build_header, + patch("requests.get") as requests_get, + patch("requests.Session.get") as session_get, + patch("requests.Session.post") as session_post, + ): + with KustoClient(kcsb) as client: + with pytest.raises(KustoClientInvalidConnectionStringException): + client.execute_query("database", "print 1") + + token_callback.assert_not_called() + cloud_info.assert_not_called() + build_header.assert_not_called() + requests_get.assert_not_called() + session_get.assert_not_called() + session_post.assert_not_called() + + +def test_cloud_login_endpoint_is_resolved_only_for_builtin_candidate(): + trusted_endpoints = KustoTrustedEndpoints() + login_endpoint = Mock(return_value=DEFAULT_PUBLIC_LOGIN_URL) + + trusted_endpoints.validate_trusted_endpoint(TRUSTED_HOST, login_endpoint) + login_endpoint.assert_called_once_with() + + login_endpoint.reset_mock() + with pytest.raises(KustoClientInvalidConnectionStringException): + trusted_endpoints.validate_trusted_endpoint(UNTRUSTED_HOST, login_endpoint) + login_endpoint.assert_not_called() + + +def test_additional_trusted_host_does_not_resolve_cloud_login_endpoint(): + trusted_endpoints = KustoTrustedEndpoints() + trusted_endpoints.add_trusted_hosts([MatchRule(".example.com", False)], False) + login_endpoint = Mock() + + trusted_endpoints.validate_trusted_endpoint(UNTRUSTED_HOST, login_endpoint) + + login_endpoint.assert_not_called() + + +def test_override_policy_preserves_additional_trusted_host_fallback(): + trusted_endpoints = KustoTrustedEndpoints() + trusted_endpoints.set_override_policy(lambda hostname: False) + trusted_endpoints.add_trusted_hosts([MatchRule(".example.com", False)], False) + login_endpoint = Mock() + + trusted_endpoints.validate_trusted_endpoint(UNTRUSTED_HOST, login_endpoint) + + login_endpoint.assert_not_called() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("authentication_method", ["user_token", "application_token", "callback"]) +async def test_async_authentication_rejects_untrusted_host_before_token_use(authentication_method): + token_callback = AsyncMock(return_value=TOKEN) + kcsb = _async_kcsb(authentication_method, token_callback) + + with ( + patch("azure.kusto.data.client_base.CloudSettings.get_cloud_info_for_cluster") as cloud_info, + patch("azure.kusto.data.security._get_header_from_dict") as build_header, + patch("requests.get") as requests_get, + patch("requests.Session.get") as session_get, + patch("azure.kusto.data.aio.client.ClientSession.post") as session_post, + ): + async with AsyncKustoClient(kcsb) as client: + with pytest.raises(KustoClientInvalidConnectionStringException): + await client.execute_query("database", "print 1") + + token_callback.assert_not_awaited() + cloud_info.assert_not_called() + build_header.assert_not_called() + requests_get.assert_not_called() + session_get.assert_not_called() + session_post.assert_not_called() \ No newline at end of file