Enable account set, list, and logout

.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

azure-cli.pyproj

@@ -22,13 +22,16 @@
     <PtvsTargetsFile>$(MSBuildExtensionsPath32)\Microsoft\VisualStudio\v$(VisualStudioVersion)\Python Tools\Microsoft.PythonTools.targets</PtvsTargetsFile>
   </PropertyGroup>
   <ItemGroup>
+    <Compile Include="azure\cli\commands\account.py" />
     <Compile Include="azure\cli\commands\login.py" />
+    <Compile Include="azure\cli\commands\logout.py" />
     <Compile Include="azure\cli\commands\storage.py" />
     <Compile Include="azure\cli\commands\__init__.py" />
     <Compile Include="azure\cli\main.py" />
     <Compile Include="azure\cli\tests\test_argparse.py">
       <SubType>Code</SubType>
     </Compile>
+    <Compile Include="azure\cli\tests\test_profile.py" />
     <Compile Include="azure\cli\_argparse.py" />
     <Compile Include="azure\cli\_logging.py" />
     <Compile Include="azure\cli\_profile.py" />

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')
+            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

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',
 ]
 

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>', _('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)

src/azure/cli/commands/login.py

@@ -1,16 +1,15 @@
-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
 
 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 <username>', _('organization Id. Microsoft Account is not yet supported.'))
 @option('--password -p <password>', _('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()
+

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 <username>', _('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)

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 <resourceGroup>', _("the resource group name"))
-@option('--subscription -s <id>', _("the subscription id"))
+@description(_('List storage accounts'))
+@option('--resource-group -g <resourceGroup>', _('the resource group name'))
+@option('--subscription -s <id>', _('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')

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()