diff --git a/vmanage/__main__.py b/vmanage/__main__.py index 3abd1e9..41f5182 100644 --- a/vmanage/__main__.py +++ b/vmanage/__main__.py @@ -23,19 +23,23 @@ def __call__(self, *args, **kwargs): class Viptela(object): - def __init__(self, host, port, username, password): + def __init__(self, host, port, username, password, tenant=None): self.host = host self.username = username self.password = password self.port = port + self.tenant = tenant self.__auth = None # use this to defer authentication until it's needed @property def auth(self): if self.__auth is None: - self.__auth = Authentication(host=self.host, port=self.port, user=self.username, - password=self.password).login() + self.__auth = Authentication(host=self.host, + port=self.port, + user=self.username, + password=self.password, + tenant=self.tenant).login() return self.__auth @@ -54,9 +58,14 @@ def auth(self): hide_input=True, help='vManage Password (env: VMANAGE_PASSWORD)', required=True) +@click.option('--tenant', + envvar='VMANAGE_TENANT', + default=None, + help='vManage Tenant Name (env: VMANAGE_TENANT)', + required=False) @click.pass_context -def vmanage(ctx, host, port, username, password): - ctx.obj = Viptela(host, port, username, password) +def vmanage(ctx, host, port, username, password, tenant=None): + ctx.obj = Viptela(host, port, username, password, tenant) vmanage.add_command(activate) diff --git a/vmanage/api/authentication.py b/vmanage/api/authentication.py index 99ab7a2..f35458b 100644 --- a/vmanage/api/authentication.py +++ b/vmanage/api/authentication.py @@ -3,6 +3,7 @@ from __future__ import (absolute_import, division, print_function) +import json import requests import urllib3 from vmanage.api.utilities import Utilities @@ -19,7 +20,7 @@ class Authentication(object): HTTP(S) Request session object will be returned. """ - def __init__(self, host=None, user=None, password=None, port=443, validate_certs=False, timeout=10): + def __init__(self, host=None, user=None, password=None, tenant=None, port=443, validate_certs=False, timeout=10): """Initialize Authentication object with session parameters. Args: @@ -37,6 +38,7 @@ def __init__(self, host=None, user=None, password=None, port=443, validate_certs self.host = host self.user = user self.password = password + self.tenant = tenant self.port = port self.timeout = timeout self.base_url = f'https://{self.host}:{self.port}/dataservice/' @@ -80,6 +82,26 @@ def login(self): response = self.session.get(url=url, timeout=self.timeout) self.session.headers['X-XSRF-TOKEN'] = response.content + if self.tenant: + api = 'tenant' + url = f'{self.base_url}{api}' + response = self.session.get(url=url, timeout=self.timeout) + tenant_list = json.loads(response.content)['data'] + tenant_pair = dict((i['name'], i['tenantId']) for i in tenant_list) + + if self.tenant in tenant_pair: + tenant_id = tenant_pair[self.tenant] + api = f'tenant/{tenant_id}/switch' + url = f'{self.base_url}{api}' + response = self.session.post(url=url, timeout=self.timeout) + + if (response.status_code != 200 or response.text.startswith('')): + raise Exception('Tenant login failed, check user credentials.') + + self.session.headers["VSessionId"] = json.loads(response.content)['VSessionId'] + else: + raise Exception('Tenant not found, check tenant name.') + except requests.exceptions.RequestException as e: raise Exception(f'Could not connect to {self.host}: {e}')