diff --git a/vmanage/api/authentication.py b/vmanage/api/authentication.py index 1f30264..226102f 100644 --- a/vmanage/api/authentication.py +++ b/vmanage/api/authentication.py @@ -5,6 +5,7 @@ import requests import urllib3 +from random import randint from vmanage.api.utilities import Utilities urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning) @@ -43,6 +44,13 @@ def __init__(self, host=None, user=None, password=None, port=443, validate_certs self.session = requests.Session() self.session.verify = validate_certs + def __enter__(self): + self.login() + return self.session + + def __exit__(self, exc_type, exc_val, exc_tb): + self.logout() + def login(self): """Executes login tasks against vManage to retrieve token(s). @@ -84,3 +92,35 @@ def login(self): raise ConnectionError(f'Could not connect to {self.host}: {e}') return self.session + + def logout(self, session=None): + """Executes a logout query against the vManage to terminate session. + + Args: + Session: Optionally pass in a session if the login function was run directly. + If no session is passed in, the logout operation is performed on self.session + + Returns: + session: Returns the same session object after the logout functionality + has been performed. + + Raises: + ConnectionError if failure to connect to vManage for logout query. + """ + try: + if not session: + session = self.session + + version = Utilities(session, self.host, self.port).get_vmanage_version() + api = f"logout?nocache={randint(1,999)}" + url = f'https://{self.host}:{self.port}/{api}' + response = session.get(url=url, timeout=self.timeout) + if response.status_code != 200: + raise ConnectionError('Logout operation failed. Recieved non-200 code from vManage') + if version >= '19.2.0': + session.headers.pop('X-XSRF-TOKEN', None) + + except requests.exceptions.RequestException as e: + raise ConnectionError(f'Could not connect to {self.host}: {e}') + + return session