From 5cf8bfdb321c07cc4ef41ac5d0176be88d5d3e11 Mon Sep 17 00:00:00 2001 From: Jiashuo Li Date: Tue, 24 Mar 2020 17:52:55 +0800 Subject: [PATCH 1/2] Extract User-Agent generation to get_az_user_agent --- .../azure/cli/core/commands/client_factory.py | 19 +- .../azure/cli/core/tests/test_util.py | 265 +++++++++--------- src/azure-cli-core/azure/cli/core/util.py | 32 ++- .../cli/command_modules/appservice/custom.py | 6 +- 4 files changed, 169 insertions(+), 153 deletions(-) diff --git a/src/azure-cli-core/azure/cli/core/commands/client_factory.py b/src/azure-cli-core/azure/cli/core/commands/client_factory.py index 6bec4216ada..478096580c8 100644 --- a/src/azure-cli-core/azure/cli/core/commands/client_factory.py +++ b/src/azure-cli-core/azure/cli/core/commands/client_factory.py @@ -3,20 +3,16 @@ # Licensed under the MIT License. See License.txt in the project root for license information. # -------------------------------------------------------------------------------------------- -import os - -from azure.cli.core import __version__ as core_version import azure.cli.core._debug as _debug from azure.cli.core.extension import EXTENSIONS_MOD_PREFIX from azure.cli.core.profiles._shared import get_client_class, SDKProfile from azure.cli.core.profiles import ResourceType, CustomResourceType, get_api_version, get_sdk +from azure.cli.core.util import get_az_user_agent from knack.log import get_logger from knack.util import CLIError logger = get_logger(__name__) -UA_AGENT = "AZURECLI/{}".format(core_version) -ENV_ADDITIONAL_USER_AGENT = 'AZURE_HTTP_USER_AGENT' def resolve_client_arg_name(operation, kwargs): @@ -82,11 +78,7 @@ def configure_common_settings(cli_ctx, client): client.config.enable_http_logger = True - client.config.add_user_agent(UA_AGENT) - try: - client.config.add_user_agent(os.environ[ENV_ADDITIONAL_USER_AGENT]) - except KeyError: - pass + client.config.add_user_agent(get_az_user_agent()) try: command_ext_name = cli_ctx.data['command_extension_name'] @@ -185,12 +177,7 @@ def get_subscription_id(cli_ctx): def _get_add_headers_callback(cli_ctx): def _add_headers(request): - agents = [request.headers['User-Agent'], UA_AGENT] - try: - agents.append(os.environ[ENV_ADDITIONAL_USER_AGENT]) - except KeyError: - pass - + agents = [request.headers['User-Agent'], get_az_user_agent()] request.headers['User-Agent'] = ' '.join(agents) try: diff --git a/src/azure-cli-core/azure/cli/core/tests/test_util.py b/src/azure-cli-core/azure/cli/core/tests/test_util.py index 3ae50ac6c64..d482413a9d3 100644 --- a/src/azure-cli-core/azure/cli/core/tests/test_util.py +++ b/src/azure-cli-core/azure/cli/core/tests/test_util.py @@ -14,7 +14,7 @@ from azure.cli.core.util import \ (get_file_json, truncate_text, shell_safe_json_parse, b64_to_hex, hash_string, random_string, open_page_in_browser, can_launch_browser, handle_exception, ConfiguredDefaultSetter, send_raw_request, - should_disable_connection_verify, parse_proxy_resource_id) + should_disable_connection_verify, parse_proxy_resource_id, get_az_user_agent) from azure.cli.core.mock import DummyCli @@ -198,6 +198,140 @@ def test_can_launch_browser(self, webbrowser_get_mock, get_platform_mock): result = can_launch_browser() self.assertFalse(result) + def test_configured_default_setter(self): + config = mock.MagicMock() + config.use_local_config = None + with ConfiguredDefaultSetter(config, True): + self.assertEqual(config.use_local_config, True) + self.assertIsNone(config.use_local_config) + + config.use_local_config = True + with ConfiguredDefaultSetter(config, False): + self.assertEqual(config.use_local_config, False) + self.assertTrue(config.use_local_config) + + @mock.patch('azure.cli.core.__version__', '7.8.9') + def test_get_az_user_agent(self): + actual = get_az_user_agent() + self.assertEqual(actual, 'AZURECLI/7.8.9') + + @mock.patch('azure.cli.core._profile.Profile.get_raw_token', autospec=True) + @mock.patch('requests.Session.send', autospec=True) + def test_send_raw_requests(self, send_mock, get_raw_token_mock): + return_val = mock.MagicMock() + return_val.is_ok = True + send_mock.return_value = return_val + get_raw_token_mock.return_value = ("Bearer", "eyJ0eXAiOiJKV1", None), None, None + + cli_ctx = DummyCli() + cli_ctx.data = { + 'command': 'rest', + 'safe_params': ['method', 'uri'] + } + test_arm_active_directory_resource_id = 'https://management.core.windows.net/' + test_arm_endpoint = 'https://management.azure.com/' + arm_resource_id = '/subscriptions/01/resourcegroups/02?api-version=2019-07-01' + full_arm_rest_url = test_arm_endpoint.rstrip('/') + arm_resource_id + test_body = '{"b1": "v1"}' + + expected_header = { + 'User-Agent': get_az_user_agent(), + 'Accept-Encoding': 'gzip, deflate', + 'Accept': '*/*', + 'Connection': 'keep-alive', + 'Content-Type': 'application/json', + 'CommandName': 'rest', + 'ParameterSetName': 'method uri', + 'Content-Length': '12' + } + expected_header_with_auth = expected_header.copy() + expected_header_with_auth['Authorization'] = 'Bearer eyJ0eXAiOiJKV1' + + # Test basic usage + # Mock Put Blob https://docs.microsoft.com/en-us/rest/api/storageservices/put-blob + # Authenticate with service SAS https://docs.microsoft.com/en-us/rest/api/storageservices/create-service-sas + sas_token = ['sv=2019-02-02', '{"srt": "s"}', "{'ss': 'bf'}"] + send_raw_request(cli_ctx, 'PUT', 'https://myaccount.blob.core.windows.net/mycontainer/myblob?timeout=30', + uri_parameters=sas_token, body=test_body, + generated_client_request_id_name=None) + + get_raw_token_mock.assert_not_called() + request = send_mock.call_args.args[1] + self.assertEqual(request.method, 'PUT') + self.assertEqual(request.url, 'https://myaccount.blob.core.windows.net/mycontainer/myblob?timeout=30&sv=2019-02-02&srt=s&ss=bf') + self.assertEqual(request.body, '{"b1": "v1"}') + # Verify no Authorization header + self.assertDictEqual(dict(request.headers), expected_header) + self.assertEqual(send_mock.call_args.kwargs["verify"], not should_disable_connection_verify()) + + # Test Authorization header is skipped + send_raw_request(cli_ctx, 'GET', full_arm_rest_url, body=test_body, skip_authorization_header=True, + generated_client_request_id_name=None) + + get_raw_token_mock.assert_not_called() + request = send_mock.call_args.args[1] + self.assertDictEqual(dict(request.headers), expected_header) + + # Test Authorization header is already provided + send_raw_request(cli_ctx, 'GET', full_arm_rest_url, + body=test_body, headers={'Authorization=Basic ABCDE'}, + generated_client_request_id_name=None) + + get_raw_token_mock.assert_not_called() + request = send_mock.call_args.args[1] + self.assertDictEqual(dict(request.headers), {**expected_header, 'Authorization': 'Basic ABCDE'}) + + # Test Authorization header is auto appended + send_raw_request(cli_ctx, 'GET', full_arm_rest_url, + body=test_body, + generated_client_request_id_name=None) + + get_raw_token_mock.assert_called_with(mock.ANY, test_arm_active_directory_resource_id) + request = send_mock.call_args.args[1] + self.assertDictEqual(dict(request.headers), expected_header_with_auth) + + # Test ARM resource ID /subscriptions/01/resourcegroups/02?api-version=2019-07-01 + send_raw_request(cli_ctx, 'GET', arm_resource_id, body=test_body, + generated_client_request_id_name=None) + + get_raw_token_mock.assert_called_with(mock.ANY, test_arm_active_directory_resource_id) + request = send_mock.call_args.args[1] + self.assertEqual(request.url, 'https://management.azure.com/subscriptions/01/resourcegroups/02?api-version=2019-07-01') + self.assertDictEqual(dict(request.headers), expected_header_with_auth) + + # Test full ARM URL https://management.azure.com/subscriptions/01/resourcegroups/02?api-version=2019-07-01 + send_raw_request(cli_ctx, 'GET', full_arm_rest_url) + + get_raw_token_mock.assert_called_with(mock.ANY, test_arm_active_directory_resource_id) + request = send_mock.call_args.args[1] + self.assertEqual(request.url, 'https://management.azure.com/subscriptions/01/resourcegroups/02?api-version=2019-07-01') + + # Test full ARM URL with port https://management.azure.com:443/subscriptions/01/resourcegroups/02?api-version=2019-07-01 + test_arm_endpoint_with_port = 'https://management.azure.com:443/' + full_arm_rest_url_with_port = test_arm_endpoint_with_port.rstrip('/') + arm_resource_id + send_raw_request(cli_ctx, 'GET', full_arm_rest_url_with_port) + + get_raw_token_mock.assert_called_with(mock.ANY, test_arm_active_directory_resource_id) + request = send_mock.call_args.args[1] + self.assertEqual(request.url, 'https://management.azure.com:443/subscriptions/01/resourcegroups/02?api-version=2019-07-01') + + # Test non-ARM API, such as MS Graph API https://graph.microsoft.com/beta/appRoleAssignments/01 + send_raw_request(cli_ctx, 'PATCH', 'https://graph.microsoft.com/beta/appRoleAssignments/01', + body=test_body, generated_client_request_id_name=None) + + get_raw_token_mock.assert_called_with(mock.ANY, 'https://graph.microsoft.com/') + request = send_mock.call_args.args[1] + self.assertEqual(request.method, 'PATCH') + self.assertEqual(request.url, 'https://graph.microsoft.com/beta/appRoleAssignments/01') + + # Test custom case-insensitive User-Agent + with mock.patch.dict('os.environ', {'AZURE_HTTP_USER_AGENT': "env-ua"}): + send_raw_request(cli_ctx, 'GET', full_arm_rest_url, headers={'user-agent=ARG-UA'}) + + get_raw_token_mock.assert_called_with(mock.ANY, test_arm_active_directory_resource_id) + request = send_mock.call_args.args[1] + self.assertEqual(request.headers['User-Agent'], get_az_user_agent() + ' env-ua ARG-UA') + class TestBase64ToHex(unittest.TestCase): @@ -333,135 +467,6 @@ def test_handle_exception_httpoperationerror_no_response_text(self, mock_logger_ self.assertEqual(mock.call(mock_http_error), mock_logger_error.call_args) self.assertEqual(ex_result, 1) - def test_configured_default_setter(self): - config = mock.MagicMock() - config.use_local_config = None - with ConfiguredDefaultSetter(config, True): - self.assertEqual(config.use_local_config, True) - self.assertIsNone(config.use_local_config) - - config.use_local_config = True - with ConfiguredDefaultSetter(config, False): - self.assertEqual(config.use_local_config, False) - self.assertTrue(config.use_local_config) - - @mock.patch('azure.cli.core._profile.Profile.get_raw_token', autospec=True) - @mock.patch('requests.Session.send', autospec=True) - def test_send_raw_requests(self, send_mock, get_raw_token_mock): - from azure.cli.core.commands.client_factory import UA_AGENT - return_val = mock.MagicMock() - return_val.is_ok = True - send_mock.return_value = return_val - get_raw_token_mock.return_value = ("Bearer", "eyJ0eXAiOiJKV1", None), None, None - - cli_ctx = DummyCli() - cli_ctx.data = { - 'command': 'rest', - 'safe_params': ['method', 'uri'] - } - test_arm_active_directory_resource_id = 'https://management.core.windows.net/' - test_arm_endpoint = 'https://management.azure.com/' - arm_resource_id = '/subscriptions/01/resourcegroups/02?api-version=2019-07-01' - full_arm_rest_url = test_arm_endpoint.rstrip('/') + arm_resource_id - test_body = '{"b1": "v1"}' - - expected_header = { - 'User-Agent': UA_AGENT, - 'Accept-Encoding': 'gzip, deflate', - 'Accept': '*/*', - 'Connection': 'keep-alive', - 'Content-Type': 'application/json', - 'CommandName': 'rest', - 'ParameterSetName': 'method uri', - 'Content-Length': '12' - } - expected_header_with_auth = expected_header.copy() - expected_header_with_auth['Authorization'] = 'Bearer eyJ0eXAiOiJKV1' - - # Test basic usage - # Mock Put Blob https://docs.microsoft.com/en-us/rest/api/storageservices/put-blob - # Authenticate with service SAS https://docs.microsoft.com/en-us/rest/api/storageservices/create-service-sas - sas_token = ['sv=2019-02-02', '{"srt": "s"}', "{'ss': 'bf'}"] - send_raw_request(cli_ctx, 'PUT', 'https://myaccount.blob.core.windows.net/mycontainer/myblob?timeout=30', - uri_parameters=sas_token, body=test_body, - generated_client_request_id_name=None) - - get_raw_token_mock.assert_not_called() - request = send_mock.call_args.args[1] - self.assertEqual(request.method, 'PUT') - self.assertEqual(request.url, 'https://myaccount.blob.core.windows.net/mycontainer/myblob?timeout=30&sv=2019-02-02&srt=s&ss=bf') - self.assertEqual(request.body, '{"b1": "v1"}') - # Verify no Authorization header - self.assertDictEqual(dict(request.headers), expected_header) - self.assertEqual(send_mock.call_args.kwargs["verify"], not should_disable_connection_verify()) - - # Test Authorization header is skipped - send_raw_request(cli_ctx, 'GET', full_arm_rest_url, body=test_body, skip_authorization_header=True, - generated_client_request_id_name=None) - - get_raw_token_mock.assert_not_called() - request = send_mock.call_args.args[1] - self.assertDictEqual(dict(request.headers), expected_header) - - # Test Authorization header is already provided - send_raw_request(cli_ctx, 'GET', full_arm_rest_url, - body=test_body, headers={'Authorization=Basic ABCDE'}, - generated_client_request_id_name=None) - - get_raw_token_mock.assert_not_called() - request = send_mock.call_args.args[1] - self.assertDictEqual(dict(request.headers), {**expected_header, 'Authorization': 'Basic ABCDE'}) - - # Test Authorization header is auto appended - send_raw_request(cli_ctx, 'GET', full_arm_rest_url, - body=test_body, - generated_client_request_id_name=None) - - get_raw_token_mock.assert_called_with(mock.ANY, test_arm_active_directory_resource_id) - request = send_mock.call_args.args[1] - self.assertDictEqual(dict(request.headers), expected_header_with_auth) - - # Test ARM resource ID /subscriptions/01/resourcegroups/02?api-version=2019-07-01 - send_raw_request(cli_ctx, 'GET', arm_resource_id, body=test_body, - generated_client_request_id_name=None) - - get_raw_token_mock.assert_called_with(mock.ANY, test_arm_active_directory_resource_id) - request = send_mock.call_args.args[1] - self.assertEqual(request.url, 'https://management.azure.com/subscriptions/01/resourcegroups/02?api-version=2019-07-01') - self.assertDictEqual(dict(request.headers), expected_header_with_auth) - - # Test full ARM URL https://management.azure.com/subscriptions/01/resourcegroups/02?api-version=2019-07-01 - send_raw_request(cli_ctx, 'GET', full_arm_rest_url) - - get_raw_token_mock.assert_called_with(mock.ANY, test_arm_active_directory_resource_id) - request = send_mock.call_args.args[1] - self.assertEqual(request.url, 'https://management.azure.com/subscriptions/01/resourcegroups/02?api-version=2019-07-01') - - # Test full ARM URL with port https://management.azure.com:443/subscriptions/01/resourcegroups/02?api-version=2019-07-01 - test_arm_endpoint_with_port = 'https://management.azure.com:443/' - full_arm_rest_url_with_port = test_arm_endpoint_with_port.rstrip('/') + arm_resource_id - send_raw_request(cli_ctx, 'GET', full_arm_rest_url_with_port) - - get_raw_token_mock.assert_called_with(mock.ANY, test_arm_active_directory_resource_id) - request = send_mock.call_args.args[1] - self.assertEqual(request.url, 'https://management.azure.com:443/subscriptions/01/resourcegroups/02?api-version=2019-07-01') - - # Test non-ARM API, such as MS Graph API https://graph.microsoft.com/beta/appRoleAssignments/01 - send_raw_request(cli_ctx, 'PATCH', 'https://graph.microsoft.com/beta/appRoleAssignments/01', - body=test_body, generated_client_request_id_name=None) - - get_raw_token_mock.assert_called_with(mock.ANY, 'https://graph.microsoft.com/') - request = send_mock.call_args.args[1] - self.assertEqual(request.method, 'PATCH') - self.assertEqual(request.url, 'https://graph.microsoft.com/beta/appRoleAssignments/01') - - # Test custom case-insensitive User-Agent - send_raw_request(cli_ctx, 'GET', full_arm_rest_url, headers={'user-agent=MY UA'}) - - get_raw_token_mock.assert_called_with(mock.ANY, test_arm_active_directory_resource_id) - request = send_mock.call_args.args[1] - self.assertEqual(request.headers['User-Agent'], 'MY UA') - @staticmethod def _get_mock_HttpOperationError(response_text): from msrest.exceptions import HttpOperationError diff --git a/src/azure-cli-core/azure/cli/core/util.py b/src/azure-cli-core/azure/cli/core/util.py index c36f50aa8f0..0aae22c4ce9 100644 --- a/src/azure-cli-core/azure/cli/core/util.py +++ b/src/azure-cli-core/azure/cli/core/util.py @@ -543,7 +543,6 @@ def send_raw_request(cli_ctx, method, uri, headers=None, uri_parameters=None, # import uuid from requests import Session, Request from requests.structures import CaseInsensitiveDict - from azure.cli.core.commands.client_factory import UA_AGENT result = CaseInsensitiveDict() for s in headers or []: @@ -559,9 +558,20 @@ def send_raw_request(cli_ctx, method, uri, headers=None, uri_parameters=None, # if 'Authorization' in headers: skip_authorization_header = True - # Allow the user to provide custom User-Agent - if 'User-Agent' not in headers: - headers['User-Agent'] = UA_AGENT + # Handle User-Agent + agents = [get_az_user_agent()] + + # Borrow AZURE_HTTP_USER_AGENT from msrest + # https://github.com/Azure/msrest-for-python/blob/4cc8bc84e96036f03b34716466230fb257e27b36/msrest/pipeline/universal.py#L70 + _ENV_ADDITIONAL_USER_AGENT = 'AZURE_HTTP_USER_AGENT' + import os + if _ENV_ADDITIONAL_USER_AGENT in os.environ: + agents.append(os.environ[_ENV_ADDITIONAL_USER_AGENT]) + + # Custom User-Agent provided as command argument + if 'User-Agent' in headers: + agents.append(headers['User-Agent']) + headers['User-Agent'] = ' '.join(agents) if generated_client_request_id_name: headers[generated_client_request_id_name] = str(uuid.uuid4()) @@ -787,3 +797,17 @@ def parse_proxy_resource_id(rid): result.pop('children', None) return {key: value for key, value in result.items() if value is not None} return None + + +def get_az_user_agent(): + # Dynamically load the core version + from azure.cli.core import __version__ as core_version + + agents = ["AZURECLI/{}".format(core_version)] + + # msrest already has this + # https://github.com/Azure/msrest-for-python/blob/4cc8bc84e96036f03b34716466230fb257e27b36/msrest/pipeline/universal.py#L70 + # if ENV_ADDITIONAL_USER_AGENT in os.environ: + # agents.append(os.environ[ENV_ADDITIONAL_USER_AGENT]) + + return ' '.join(agents) diff --git a/src/azure-cli/azure/cli/command_modules/appservice/custom.py b/src/azure-cli/azure/cli/command_modules/appservice/custom.py index b044cb8d028..cb00b5a512d 100644 --- a/src/azure-cli/azure/cli/command_modules/appservice/custom.py +++ b/src/azure-cli/azure/cli/command_modules/appservice/custom.py @@ -43,7 +43,7 @@ from azure.cli.core.commands import LongRunningOperation from azure.cli.core.util import in_cloud_console, shell_safe_json_parse, open_page_in_browser, get_json_object, \ ConfiguredDefaultSetter, sdk_no_wait -from azure.cli.core.commands.client_factory import UA_AGENT +from azure.cli.core.util import get_az_user_agent from azure.cli.core.profiles import ResourceType from .tunnel import TunnelServer @@ -358,7 +358,7 @@ def enable_zip_deploy(cmd, resource_group_name, name, src, timeout=None, slot=No headers = authorization headers['Content-Type'] = 'application/octet-stream' headers['Cache-Control'] = 'no-cache' - headers['User-Agent'] = UA_AGENT + headers['User-Agent'] = get_az_user_agent() import requests import os @@ -827,7 +827,7 @@ def _get_app_settings_from_scm(cmd, resource_group_name, name, slot=None): headers = { 'Content-Type': 'application/octet-stream', 'Cache-Control': 'no-cache', - 'User-Agent': UA_AGENT + 'User-Agent': get_az_user_agent() } import requests From 183935906ff76e3801f6b62f87d5163d4b35ccdf Mon Sep 17 00:00:00 2001 From: Jiashuo Li Date: Tue, 24 Mar 2020 22:51:41 +0800 Subject: [PATCH 2/2] patch os.environ --- src/azure-cli-core/azure/cli/core/tests/test_util.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/src/azure-cli-core/azure/cli/core/tests/test_util.py b/src/azure-cli-core/azure/cli/core/tests/test_util.py index d482413a9d3..2b6536606fa 100644 --- a/src/azure-cli-core/azure/cli/core/tests/test_util.py +++ b/src/azure-cli-core/azure/cli/core/tests/test_util.py @@ -5,6 +5,7 @@ # pylint: disable=line-too-long from collections import namedtuple +import os import sys import unittest import mock @@ -215,9 +216,13 @@ def test_get_az_user_agent(self): actual = get_az_user_agent() self.assertEqual(actual, 'AZURECLI/7.8.9') + @mock.patch.dict('os.environ') @mock.patch('azure.cli.core._profile.Profile.get_raw_token', autospec=True) @mock.patch('requests.Session.send', autospec=True) def test_send_raw_requests(self, send_mock, get_raw_token_mock): + if 'AZURE_HTTP_USER_AGENT' in os.environ: + del os.environ['AZURE_HTTP_USER_AGENT'] # Clear env var possibly added by DevOps + return_val = mock.MagicMock() return_val.is_ok = True send_mock.return_value = return_val