From 0d76d4ca7ec9953b80b7745374fafcabece23ce6 Mon Sep 17 00:00:00 2001
From: Josh Pearson <josh_jip_pearson@hotmail.com>
Date: Thu, 4 Jan 2024 23:48:56 +0000
Subject: [PATCH 1/3] Add logout functionality to API object with the ability
 to use a context manager - Created logout function which sends a request the
 the logout endpoint - Created __enter__ and __exit__ methods which allow for
 use of a context manager

---
 vmanage/api/authentication.py | 41 +++++++++++++++++++++++++++++++++++
 1 file changed, 41 insertions(+)

diff --git a/vmanage/api/authentication.py b/vmanage/api/authentication.py
index 1f30264..dd08bba 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)
@@ -42,6 +43,13 @@ def __init__(self, host=None, user=None, password=None, port=443, validate_certs
         self.base_url = f'https://{self.host}:{self.port}/dataservice/'
         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,36 @@ 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}')
+
+        print("executed logout")
+        return session
\ No newline at end of file

From d31235acca9a6ef26a08cc6d34a9710554d96398 Mon Sep 17 00:00:00 2001
From: Josh Pearson <josh_jip_pearson@hotmail.com>
Date: Fri, 5 Jan 2024 00:02:48 +0000
Subject: [PATCH 2/3] Fix some formatting with make format

---
 vmanage/api/authentication.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/vmanage/api/authentication.py b/vmanage/api/authentication.py
index dd08bba..e317b20 100644
--- a/vmanage/api/authentication.py
+++ b/vmanage/api/authentication.py
@@ -43,7 +43,7 @@ def __init__(self, host=None, user=None, password=None, port=443, validate_certs
         self.base_url = f'https://{self.host}:{self.port}/dataservice/'
         self.session = requests.Session()
         self.session.verify = validate_certs
-    
+
     def __enter__(self):
         self.login()
         return self.session
@@ -92,7 +92,7 @@ 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.
 
@@ -124,4 +124,4 @@ def logout(self, session=None):
             raise ConnectionError(f'Could not connect to {self.host}: {e}')
 
         print("executed logout")
-        return session
\ No newline at end of file
+        return session

From 1b1e5dcc4c28401508600b2494bceb4195ae6175 Mon Sep 17 00:00:00 2001
From: Josh Pearson <josh_jip_pearson@hotmail.com>
Date: Fri, 5 Jan 2024 00:10:23 +0000
Subject: [PATCH 3/3] Remove test message

---
 vmanage/api/authentication.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/vmanage/api/authentication.py b/vmanage/api/authentication.py
index e317b20..226102f 100644
--- a/vmanage/api/authentication.py
+++ b/vmanage/api/authentication.py
@@ -123,5 +123,4 @@ def logout(self, session=None):
         except requests.exceptions.RequestException as e:
             raise ConnectionError(f'Could not connect to {self.host}: {e}')
 
-        print("executed logout")
         return session