create http client class

.gitignore

@@ -91,6 +91,7 @@ venv/
 ENV/
 env.bak/
 venv.bak/
+.vscode
 
 # Spyder project settings
 .spyderproject

tests/functional_test_http_client.py

@@ -0,0 +1,57 @@
+# vim: set filetype=python ts=4 sw=4
+# -*- coding: utf-8 -*-
+import unittest
+
+from tokendito.http_client import HTTPClient
+from tokendito import __title__
+from tokendito import __version__
+
+class TestHTTPClientFunctional(unittest.TestCase):
+    def setUp(self):
+        # Initialize the HTTPClient with a test user-agent
+        self.client = HTTPClient()
+        self.client.session.headers.update({"User-Agent": f"{__title__}/{__version__}"})
+
+
+    def test_get_request(self):
+        # Make a GET request to the /get endpoint of httpbin which reflects the sent request data
+        response = self.client.get("https://httpbin.org/get")
+        json_data = response.json()
+
+        # Assert that the request was successful and the returned User-Agent matches the one we set
+        self.assertEqual(response.status_code, 200)
+        self.assertEqual(json_data["headers"]["User-Agent"], f"{__title__}/{__version__}")
+
+    def test_post_request(self):
+        # Make a POST request to the /post endpoint of httpbin with sample data
+        response = self.client.post("https://httpbin.org/post", json={"key": "value"})
+        json_data = response.json()
+
+        # Assert that the request was successful and the returned json data matches the data we sent
+        self.assertEqual(response.status_code, 200)
+        self.assertEqual(json_data["json"], {"key": "value"})
+
+    def test_set_cookies(self):
+        # Set a test cookie for the client
+        self.client.set_cookies({"test_cookie": "cookie_value"})
+
+        # Make a request to the /cookies endpoint of httpbin which returns set cookies
+        response = self.client.get("https://httpbin.org/cookies")
+        json_data = response.json()
+
+        # Assert that the cookie we set is correctly returned by the server
+        self.assertEqual(json_data["cookies"], {"test_cookie": "cookie_value"})
+
+    def test_custom_header(self):
+        # Make a GET request with a custom header
+        response = self.client.get(
+            "https://httpbin.org/get", headers={"X-Test-Header": "TestValue"}
+        )
+        json_data = response.json()
+
+        # Assert that the custom header was correctly sent
+        self.assertEqual(json_data["headers"]["X-Test-Header"], "TestValue")
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/test_http_client.py

@@ -0,0 +1,75 @@
+# vim: set filetype=python ts=4 sw=4
+# -*- coding: utf-8 -*-
+import unittest
+from unittest.mock import Mock, patch
+
+import requests
+from tokendito.http_client import HTTPClient
+from tokendito import __title__
+from tokendito import __version__
+
+class TestHTTPClient(unittest.TestCase):
+    """
+    Unit test class for the HTTPClient.
+    """
+
+    def setUp(self):
+        """
+        Setup method called before each test. Initializes an HTTPClient instance.
+        """
+        # Initializing HTTPClient instance without the 'user_agent' parameter
+        self.client = HTTPClient()
+
+    def test_init(self):
+        """
+        Test the initialization of the HTTPClient instance and its properties.
+        """
+        # Check if the session property of the client is an instance of requests.Session
+        self.assertIsInstance(self.client.session, requests.Session)
+
+        # Check if the User-Agent header was set correctly during initialization
+        # Replace 'ExpectedUserAgentStringHere' with the actual User-Agent string from your HTTPClient class
+        expected_user_agent = f"{__title__}/{__version__}"
+        self.assertEqual(self.client.session.headers["User-Agent"], expected_user_agent)
+
+    def test_set_cookies(self):
+        """
+        Test setting cookies for the HTTPClient session.
+        """
+        cookies = {"test_cookie": "cookie_value"}
+        self.client.set_cookies(cookies)
+        # Check if the provided cookie is set correctly in the session
+        self.assertEqual(self.client.session.cookies.get_dict(), cookies)
+
+    @patch("requests.Session.get")
+    def test_get(self, mock_get):
+        """
+        Test the GET request method of the HTTPClient. Uses mocking to simulate an HTTP response.
+        """
+        mock_resp = Mock()
+        mock_resp.status_code = 200
+        mock_resp.text = "OK"
+        mock_get.return_value = mock_resp
+
+        response = self.client.get("http://test.com")
+        # Check if the response status code and text match the expected values
+        self.assertEqual(response.status_code, 200)
+        self.assertEqual(response.text, "OK")
+
+    @patch("requests.Session.post")
+    def test_post(self, mock_post):
+        """
+        Test the POST request method of the HTTPClient. Uses mocking to simulate an HTTP response.
+        """
+        mock_resp = Mock()
+        mock_resp.status_code = 201
+        mock_resp.text = "Created"
+        mock_post.return_value = mock_resp
+
+        response = self.client.post("http://test.com", json={"key": "value"})
+        # Check if the response status code and text match the expected values
+        self.assertEqual(response.status_code, 201)
+        self.assertEqual(response.text, "Created")
+
+if __name__ == "__main__":
+    unittest.main()

tokendito/__init__.py

@@ -1,7 +1,7 @@
 # vim: set filetype=python ts=4 sw=4
 # -*- coding: utf-8 -*-
 """Tokendito module initialization."""
-__version__ = "2.1.2"
+__version__ = "2.1.3"
 __title__ = "tokendito"
 __description__ = "Get AWS STS tokens from Okta SSO"
 __long_description_content_type__ = "text/markdown"

tokendito/aws.py

@@ -16,6 +16,7 @@
 from botocore.client import Config as BotocoreConfig
 from botocore.exceptions import ClientError
 import botocore.session
+from http_client import HTTP_client
 from tokendito import okta
 from tokendito import user
 
@@ -51,19 +52,21 @@ def authenticate_to_roles(urls, cookies=None):
 
     :param urls: list of tuples or tuple, with tiles info
     :param cookies: html cookies
+    :param user_agent: optional user agent string
     :return: response text
 
     """
+    if cookies:
+        HTTP_client.set_cookies(cookies)  # Set cookies if provided
+
     url_list = [urls] if isinstance(urls, tuple) else urls
     responses = []
     tile_count = len(url_list)
-    plural = ""
-    if tile_count > 1:
-        plural = "s"
+    plural = "s" if tile_count > 1 else ""
 
     logger.info(f"Discovering roles in {tile_count} tile{plural}.")
     for url, label in url_list:
-        response = user.request_wrapper("GET", url, cookies=cookies)
+        response = HTTP_client.get(url)  # Use the HTTPClient's get method
         saml_response_string = response.text
 
         saml_xml = okta.extract_saml_response(saml_response_string)

tokendito/http_client.py

@@ -0,0 +1,97 @@
+# vim: set filetype=python ts=4 sw=4
+# -*- coding: utf-8 -*-
+import logging
+
+import requests
+from tokendito import __title__
+from tokendito import __version__
+
+logger = logging.getLogger(__name__)
+
+
+class HTTPClient:
+    def __init__(self):
+        """
+        Initialize the HTTPClient with a session object.
+        """
+
+        # Set the default User-Agent using the module's title and version
+        user_agent = f"{__title__}/{__version__}"
+
+        # Create a session object that will be used for making HTTP calls
+        self.session = requests.Session()
+
+        # Set the User-Agent header for the session
+        self.session.headers.update({"User-Agent": user_agent})
+
+    def set_cookies(self, cookies):
+        """
+        Update the session with additional cookies.
+
+        Args:
+            cookies (dict): A dictionary of cookies to be added to the session.
+        """
+        # Update the session cookies with the provided cookies
+        self.session.cookies.update(cookies)
+
+    def get(self, url, params=None, headers=None):
+        """
+        Perform a GET request using the session.
+
+        Args:
+            url (str): The URL to send the GET request to.
+            params (dict, optional): The query parameters to be added to the URL.
+            headers (dict, optional): Additional headers to include in the request.
+
+        Returns:
+            Response: The response object resulting from the GET request.
+        """
+        try:
+            logger.debug(f"Sending cookies: {self.session.cookies}")
+            logger.debug(f"Sending headers: {self.session.headers}")
+            response = self.session.get(url, params=params, headers=headers)
+            response.raise_for_status()
+            logger.debug(f"Received response from {url}: {response.text}")
+            return response
+        except requests.RequestException as e:
+            logger.error(f"Error during GET request to {url}. Error: {e}")
+            logger.debug(f"Response Headers: {response.headers}")
+            logger.debug(f"Response Content: {response.content}")
+            raise
+
+    def post(self, url, data=None, json=None, headers=None):
+        """
+        Perform a POST request using the session.
+
+        Args:
+            url (str): The URL to send the POST request to.
+            data (dict, optional): The data to send in the body of the POST request.
+            json (dict, optional): JSON data to send in the body of the POST request.
+            headers (dict, optional): Additional headers to include in the request.
+
+        Returns:
+            Response: The response object resulting from the POST request.
+        """
+        try:
+            response = self.session.post(url, data=data, json=json, headers=headers)
+            response.raise_for_status()
+            return response
+        except requests.RequestException as e:
+            logger.error(f"Error during POST request to {url}. Error: {e}")
+            raise
+
+    def reset(self):
+        """
+        Resets the session object to its initial state.
+        """
+        user_agent = f"{__title__}/{__version__}"  # Set the default User-Agent
+        self.session.cookies.clear()  # Clear any cookies that may have been set
+        self.session.headers = (
+            requests.utils.default_headers()
+        )  # Clear any headers that may have been set
+        self.session.headers.update(
+            {"User-Agent": user_agent}
+        )  # Set the User-Agent header for the session
+
+
+HTTP_client = HTTPClient()