From 9319be15b65d9193d12aa686bece2ceb1cd423be Mon Sep 17 00:00:00 2001 From: yugangw-msft Date: Fri, 19 Feb 2016 16:02:14 -0800 Subject: [PATCH] Enable account set, list, and logout --- .travis.yml | 2 + azure-cli.pyproj | 3 + src/azure/cli/_profile.py | 105 +++++++++++++++++--- src/azure/cli/commands/__init__.py | 2 + src/azure/cli/commands/account.py | 29 ++++++ src/azure/cli/commands/login.py | 22 ++--- src/azure/cli/commands/logout.py | 13 +++ src/azure/cli/commands/storage.py | 13 ++- src/azure/cli/tests/test_argparse.py | 12 +-- src/azure/cli/tests/test_profile.py | 141 +++++++++++++++++++++++++++ 10 files changed, 300 insertions(+), 42 deletions(-) create mode 100644 src/azure/cli/commands/account.py create mode 100644 src/azure/cli/commands/logout.py create mode 100644 src/azure/cli/tests/test_profile.py diff --git a/.travis.yml b/.travis.yml index 06fc3541497..adbf848c3ed 100644 --- a/.travis.yml +++ b/.travis.yml @@ -3,6 +3,8 @@ language: python python: - "2.7" - "3.5" +install: + - pip install azure==2.0.0a1 script: - export PYTHONPATH=$PATHONPATH:./src - python -m unittest discover -s src/azure/cli/tests \ No newline at end of file diff --git a/azure-cli.pyproj b/azure-cli.pyproj index f0c24e1b38b..ddc1c1edb89 100644 --- a/azure-cli.pyproj +++ b/azure-cli.pyproj @@ -22,13 +22,16 @@ $(MSBuildExtensionsPath32)\Microsoft\VisualStudio\v$(VisualStudioVersion)\Python Tools\Microsoft.PythonTools.targets + + Code + diff --git a/src/azure/cli/_profile.py b/src/azure/cli/_profile.py index 50190a5a6c6..1f663ba581b 100644 --- a/src/azure/cli/_profile.py +++ b/src/azure/cli/_profile.py @@ -1,22 +1,99 @@ from msrest.authentication import BasicTokenAuthentication - from .main import CONFIG +import collections class Profile(object): - def update(self, subscriptions, access_token): - subscriptions[0]['active'] = True - CONFIG['subscriptions'] = subscriptions - CONFIG['access_token'] = access_token + def __init__(self, storage=CONFIG): + self._storage = storage + + @staticmethod + def normalize_properties(user, subscriptions): + consolidated = [] + for s in subscriptions: + consolidated.append({ + 'id': s.id.split('/')[-1], + 'name': s.display_name, + 'state': s.state, + 'user': user, + 'active': False + }) + return consolidated + + def set_subscriptions(self, new_subscriptions, access_token): + existing_ones = self.load_subscriptions() + active_one = next((x for x in existing_ones if x['active']), None) + active_subscription_id = active_one['id'] if active_one else None + + #merge with existing ones + dic = collections.OrderedDict((x['id'], x) for x in existing_ones) + dic.update((x['id'], x) for x in new_subscriptions) + subscriptions = list(dic.values()) + + if active_one: + new_active_one = next( + (x for x in new_subscriptions if x['id'] == active_subscription_id), None) - def get_credentials(self): - subscriptions = CONFIG['subscriptions'] - sub = [x for x in subscriptions if x['active'] == True ] - if not sub and subscriptions: - sub = subscriptions + for s in subscriptions: + s['active'] = False - if sub: - return (BasicTokenAuthentication({ 'access_token': CONFIG['access_token']}), - sub[0]['id'] ) + if new_active_one: + new_active_one['active'] = True + else: + new_subscriptions[0]['active'] = True else: - raise ValueError('you need to login to') \ No newline at end of file + new_subscriptions[0]['active'] = True + + #before adal/python is available, persist tokens with other profile info + for s in new_subscriptions: + s['access_token'] = access_token + + self._save_subscriptions(subscriptions) + + def get_login_credentials(self): + subscriptions = self.load_subscriptions() + if not subscriptions: + raise ValueError('Please run login to setup account.') + + active = [x for x in subscriptions if x['active']] + if len(active) != 1: + raise ValueError('Please run "account set" to select active account.') + + return BasicTokenAuthentication( + {'access_token': active[0]['access_token']}), active[0]['id'] + + def set_active_subscription(self, subscription_id_or_name): + subscriptions = self.load_subscriptions() + + subscription_id_or_name = subscription_id_or_name.lower() + result = [x for x in subscriptions + if subscription_id_or_name == x['id'].lower() or + subscription_id_or_name == x['name'].lower()] + + if len(result) != 1: + raise ValueError('The subscription of "{}" does not exist or has more than' + ' one match.'.format(subscription_id_or_name)) + + for s in subscriptions: + s['active'] = False + result[0]['active'] = True + + self._save_subscriptions(subscriptions) + + def logout(self, user): + subscriptions = self.load_subscriptions() + result = [x for x in subscriptions if user.lower() == x['user'].lower()] + subscriptions = [x for x in subscriptions if x not in result] + + #reset the active subscription if needed + result = [x for x in subscriptions if x['active']] + if not result and subscriptions: + subscriptions[0]['active'] = True + + self._save_subscriptions(subscriptions) + + def load_subscriptions(self): + return self._storage.get('subscriptions') or [] + + def _save_subscriptions(self, subscriptions): + self._storage['subscriptions'] = subscriptions diff --git a/src/azure/cli/commands/__init__.py b/src/azure/cli/commands/__init__.py index 598ab429816..1468fb7c7bd 100644 --- a/src/azure/cli/commands/__init__.py +++ b/src/azure/cli/commands/__init__.py @@ -4,6 +4,8 @@ # TODO: Alternatively, simply scan the directory for all modules COMMAND_MODULES = [ 'login', + 'logout', + 'account', 'storage', ] diff --git a/src/azure/cli/commands/account.py b/src/azure/cli/commands/account.py new file mode 100644 index 00000000000..2eb554fed68 --- /dev/null +++ b/src/azure/cli/commands/account.py @@ -0,0 +1,29 @@ +from .._profile import Profile +from .._util import TableOutput +from ..commands import command, description, option + +@command('account list') +@description(_('List the imported subscriptions.')) +def list_subscriptions(args, unexpected): + profile = Profile() + subscriptions = profile.load_subscriptions() + + with TableOutput() as to: + for subscription in subscriptions: + to.cell('Name', subscription['name']) + to.cell('Active', bool(subscription['active'])) + to.cell('User', subscription['user']) + to.cell('Subscription Id', subscription['id']) + to.cell('State', subscription['state']) + to.end_row() + +@command('account set') +@description(_('Set the current subscription')) +@option('--subscription-id -n ', _('Subscription Id, unique name also works.')) +def set_active_subscription(args, unexpected): + id = args.get('subscription-id') + if not id: + raise ValueError(_('Please provide subscription id or unique name.')) + + profile = Profile() + profile.set_active_subscription(id) diff --git a/src/azure/cli/commands/login.py b/src/azure/cli/commands/login.py index f91a60bd323..447096e5437 100644 --- a/src/azure/cli/commands/login.py +++ b/src/azure/cli/commands/login.py @@ -1,8 +1,7 @@ -from azure.mgmt.resource.subscriptions import SubscriptionClient, \ +from msrestazure.azure_active_directory import UserPassCredentials +from azure.mgmt.resource.subscriptions import SubscriptionClient, \ SubscriptionClientConfiguration -from msrestazure.azure_active_directory import UserPassCredentials -from .._logging import logger from .._profile import Profile from .._util import TableOutput from ..commands import command, description, option @@ -10,7 +9,7 @@ CLIENT_ID = '04b07795-8ddb-461a-bbee-02f9e1bf7b46' @command('login') -@description('log in to an Azure subscription using Active Directory Organization Id') +@description(_('log in to an Azure subscription using Active Directory Organization Id')) @option('--username -u ', _('organization Id. Microsoft Account is not yet supported.')) @option('--password -p ', _('user password, will prompt if not given.')) def login(args, unexpected): @@ -26,22 +25,14 @@ def login(args, unexpected): subscriptions = client.subscriptions.list() if not subscriptions: - raise RuntimeError(_("No subscriptions found for this account")) + raise RuntimeError(_('No subscriptions found for this account.')) #keep useful properties and not json serializable - consolidated = [] - for s in subscriptions: - subscription = {}; - subscription['id'] = s.id.split('/')[-1] - subscription['name'] = s.display_name - subscription['state'] = s.state - subscription['user'] = username - consolidated.append(subscription) profile = Profile() - profile.update(consolidated, credentials.token['access_token']) + consolidated = Profile.normalize_properties(username, subscriptions) + profile.set_subscriptions(consolidated, credentials.token['access_token']) - #TODO, replace with JSON display with TableOutput() as to: for subscription in consolidated: to.cell('Name', subscription['name']) @@ -50,3 +41,4 @@ def login(args, unexpected): to.cell('Subscription Id', subscription['id']) to.cell('State', subscription['state']) to.end_row() + diff --git a/src/azure/cli/commands/logout.py b/src/azure/cli/commands/logout.py new file mode 100644 index 00000000000..341904576e7 --- /dev/null +++ b/src/azure/cli/commands/logout.py @@ -0,0 +1,13 @@ +from .._profile import Profile +from ..commands import command, description, option + +@command('logout') +@description(_('Log out from Azure subscription using Active Directory.')) +@option('--username -u ', _('User name used to log out from Azure Active Directory.')) +def logout(args, unexpected): + username = args.get('username') + if not username: + raise ValueError(_('Please provide a valid username to logout.')) + + profile = Profile() + profile.logout(username) diff --git a/src/azure/cli/commands/storage.py b/src/azure/cli/commands/storage.py index c8d4f353bc5..f235a26ca3b 100644 --- a/src/azure/cli/commands/storage.py +++ b/src/azure/cli/commands/storage.py @@ -1,22 +1,21 @@ -from ..main import CONFIG, SESSION -from .._logging import logger +from ..main import SESSION +from .._logging import logging from .._util import TableOutput from ..commands import command, description, option from .._profile import Profile @command('storage account list') -@description('List storage accounts') -@option('--resource-group -g ', _("the resource group name")) -@option('--subscription -s ', _("the subscription id")) +@description(_('List storage accounts')) +@option('--resource-group -g ', _('the resource group name')) +@option('--subscription -s ', _('the subscription id')) def list_accounts(args, unexpected): from azure.mgmt.storage import StorageManagementClient, StorageManagementClientConfiguration from azure.mgmt.storage.models import StorageAccount from msrestazure.azure_active_directory import UserPassCredentials profile = Profile() - #credentials, subscription_id = profile.get_credentials() smc = StorageManagementClient(StorageManagementClientConfiguration( - *profile.get_credentials() + *profile.get_login_credentials(), )) group = args.get('resource-group') diff --git a/src/azure/cli/tests/test_argparse.py b/src/azure/cli/tests/test_argparse.py index 76491e2e435..99f6b7fa0e3 100644 --- a/src/azure/cli/tests/test_argparse.py +++ b/src/azure/cli/tests/test_argparse.py @@ -58,12 +58,12 @@ def test_args(self): self.assertIsNone(res) res, other = p.execute('n1 -b -a x'.split()) - self.assertEquals(res.b, '-a') + self.assertEqual(res.b, '-a') self.assertSequenceEqual(res.positional, ['x']) self.assertRaises(IncorrectUsageError, lambda: res.arg) res, other = p.execute('n1 -b:-a x'.split()) - self.assertEquals(res.b, '-a') + self.assertEqual(res.b, '-a') self.assertSequenceEqual(res.positional, ['x']) self.assertRaises(IncorrectUsageError, lambda: res.arg) @@ -73,16 +73,16 @@ def test_unexpected_args(self): res, other = p.execute('n1 -b=2'.split()) self.assertFalse(res) - self.assertEquals('2', other.b) + self.assertEqual('2', other.b) res, other = p.execute('n1 -b.c.d=2'.split()) self.assertFalse(res) - self.assertEquals('2', other.b.c.d) + self.assertEqual('2', other.b.c.d) res, other = p.execute('n1 -b.c.d 2 -b.c.e:3'.split()) self.assertFalse(res) - self.assertEquals('2', other.b.c.d) - self.assertEquals('3', other.b.c.e) + self.assertEqual('2', other.b.c.d) + self.assertEqual('3', other.b.c.e) if __name__ == '__main__': unittest.main() diff --git a/src/azure/cli/tests/test_profile.py b/src/azure/cli/tests/test_profile.py new file mode 100644 index 00000000000..dd582130d80 --- /dev/null +++ b/src/azure/cli/tests/test_profile.py @@ -0,0 +1,141 @@ +import unittest +from azure.cli._profile import Profile + +class Test_Profile(unittest.TestCase): + + @classmethod + def setUpClass(cls): + cls.user1 = 'foo@foo.com' + cls.id1 = 'subscriptions/1' + cls.display_name1 = 'foo account' + cls.state1 = 'enabled' + cls.subscription1 = SubscriptionStub(cls.id1, + cls.display_name1, + cls.state1) + cls.token1 = 'token1' + + cls.user2 = 'bar@bar.com' + cls.id2 = 'subscriptions/2' + cls.display_name2 = 'bar account' + cls.state2 = 'suspended' + cls.subscription2 = SubscriptionStub(cls.id2, + cls.display_name2, + cls.state2) + cls.token2 = 'token2' + + def test_normalize(self): + consolidated = Profile.normalize_properties(self.user1, + [self.subscription1]) + self.assertEqual(consolidated[0], { + 'id': '1', + 'name': self.display_name1, + 'state': self.state1, + 'user': self.user1, + 'active': False + }) + + def test_update_add_two_different_subscriptions(self): + storage_mock = {'subscriptions': None} + profile = Profile(storage_mock) + + #add the first and verify + consolidated = Profile.normalize_properties(self.user1, + [self.subscription1]) + profile.set_subscriptions(consolidated, self.token1) + + self.assertEqual(len(storage_mock['subscriptions']), 1) + subscription1 = storage_mock['subscriptions'][0] + self.assertEqual(subscription1, { + 'id': '1', + 'name': self.display_name1, + 'state': self.state1, + 'user': self.user1, + 'access_token': self.token1, + 'active': True + }) + + #add the second and verify + consolidated = Profile.normalize_properties(self.user2, + [self.subscription2]) + profile.set_subscriptions(consolidated, self.token2) + + self.assertEqual(len(storage_mock['subscriptions']), 2) + subscription2 = storage_mock['subscriptions'][1] + self.assertEqual(subscription2, { + 'id': '2', + 'name': self.display_name2, + 'state': self.state2, + 'user': self.user2, + 'access_token': self.token2, + 'active': True + }) + + #verify the old one stays, but no longer active + self.assertEqual(storage_mock['subscriptions'][0]['name'], + subscription1['name']) + self.assertEqual(storage_mock['subscriptions'][0]['access_token'], + self.token1) + self.assertFalse(storage_mock['subscriptions'][0]['active']) + + def test_update_with_same_subscription_added_twice(self): + storage_mock = {'subscriptions': None} + profile = Profile(storage_mock) + + #add one twice and verify we will have one but with new token + consolidated = Profile.normalize_properties(self.user1, + [self.subscription1]) + profile.set_subscriptions(consolidated, self.token1) + + new_subscription1 = SubscriptionStub(self.id1, + self.display_name1, + self.state1) + consolidated = Profile.normalize_properties(self.user1, + [new_subscription1]) + profile.set_subscriptions(consolidated, self.token2) + + self.assertEqual(len(storage_mock['subscriptions']), 1) + self.assertEqual(storage_mock['subscriptions'][0]['access_token'], + self.token2) + self.assertTrue(storage_mock['subscriptions'][0]['active']) + + def test_set_active_subscription(self): + storage_mock = {'subscriptions': None} + profile = Profile(storage_mock) + + consolidated = Profile.normalize_properties(self.user1, + [self.subscription1]) + profile.set_subscriptions(consolidated, self.token1) + + consolidated = Profile.normalize_properties(self.user2, + [self.subscription2]) + profile.set_subscriptions(consolidated, self.token2) + + subscription1 = storage_mock['subscriptions'][0] + subscription2 = storage_mock['subscriptions'][1] + self.assertTrue(subscription2['active']) + + profile.set_active_subscription(subscription1['id']) + self.assertFalse(subscription2['active']) + self.assertTrue(subscription1['active']) + + def test_get_login_credentials(self): + storage_mock = {'subscriptions': None} + profile = Profile(storage_mock) + + consolidated = Profile.normalize_properties(self.user1, + [self.subscription1]) + profile.set_subscriptions(consolidated, self.token1) + cred, subscription_id = profile.get_login_credentials() + + self.assertEqual(cred.token['access_token'], self.token1) + self.assertEqual(subscription_id, '1') + + +class SubscriptionStub: + def __init__(self, id, display_name, state): + self.id = id + self.display_name = display_name + self.state = state + +if __name__ == '__main__': + unittest.main()