Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
19 changes: 3 additions & 16 deletions src/azure-cli-core/azure/cli/core/commands/client_factory.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down Expand Up @@ -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']
Expand Down Expand Up @@ -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:
Expand Down
270 changes: 140 additions & 130 deletions src/azure-cli-core/azure/cli/core/tests/test_util.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@

# pylint: disable=line-too-long
from collections import namedtuple
import os
import sys
import unittest
import mock
Expand All @@ -14,7 +15,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


Expand Down Expand Up @@ -198,6 +199,144 @@ 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')
Copy link
Member Author

@jiasli jiasli Mar 24, 2020

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

#12717 can use this test method to verify whether the installer is added to UA correctly.


@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):
Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

These test methods are moved from the wrong class TestHandleException to the right class TestUtils.

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
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')
Comment on lines +332 to +338
Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

User-Agent for az rest is tested here.



class TestBase64ToHex(unittest.TestCase):

Expand Down Expand Up @@ -333,135 +472,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
Expand Down
32 changes: 28 additions & 4 deletions src/azure-cli-core/azure/cli/core/util.py
Original file line number Diff line number Diff line change
Expand Up @@ -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 []:
Expand All @@ -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())
Expand Down Expand Up @@ -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)
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

May I know why we need the new function to handle agents? From my side, the previous UA definition should be fine, too.

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Because previously the functionality of concatenating UA is dispersed in 3 places:

  • mgmt client
  • data client
  • az rest

It is difficult to maintain. I extrated the logic to one place so that changes done here automatically applies to all 3 places, given #12717 is doing more tweaks on the UA.

Loading