diff --git a/lib/reqs.py b/lib/reqs.py index 2642e15..779e6e2 100644 --- a/lib/reqs.py +++ b/lib/reqs.py @@ -1,105 +1,84 @@ import requests import json from datetime import datetime -from typing import List, Dict, Any +from typing import List, Dict, Any, Optional + +# Create a session object to reuse TCP connections +session = requests.Session() + +def _make_api_request( + method: str, + base_url: str, + auth: str, + endpoint: str, + payload: Optional[Dict[str, Any]] = None, +) -> Any: + """ + A helper function to make API requests. + + Args: + method: The HTTP method to use (e.g., 'GET', 'POST'). + base_url: The base URL of the API. + auth: The authorization token. + endpoint: The API endpoint to call. + payload: The data to send with the request. + + Returns: + The JSON response from the API. + """ + url = f"{base_url}/{endpoint}" + headers = {"Authorization": f"Basic {auth}"} + try: + response = session.request(method, url, headers=headers, data=payload) + response.raise_for_status() + return response + except (requests.exceptions.RequestException, json.JSONDecodeError) as e: + # Log the error for debugging purposes + print(f"API request failed: {e}") + return None def get_urls(base_url: str, auth: str) -> List[Dict[str, Any]]: - """Retrieves user account streams from the API. - - Args: - base_url: The base URL of the API. - auth: The authorization token. - - Returns: - A list of user account streams. - """ - url = f"{base_url}/getUserAccounts/streams" - headers = {"Authorization": f"Basic {auth}"} - response = requests.get(url, headers=headers) - response.raise_for_status() - return response.json() + """Retrieves user account streams from the API.""" + response = _make_api_request("GET", base_url, auth, "getUserAccounts/streams") + return response.json() if response else [] def get_user_accounts(base_url: str, auth: str) -> List[Dict[str, Any]]: - """Retrieves user accounts from the API. + """Retrieves user accounts from the API.""" + response = _make_api_request("GET", base_url, auth, "getUserAccounts") + if not response: + return [] - Args: - base_url: The base URL of the API. - auth: The authorization token. - - Returns: - A list of user accounts with 'expiaryDate_rendered' added. - """ - url = f"{base_url}/getUserAccounts" - headers = {"Authorization": f"Basic {auth}"} - response = requests.get(url, headers=headers) - response.raise_for_status() accounts = response.json() - for account in accounts: account["expiaryDate_rendered"] = datetime.utcfromtimestamp( account["expiaryDate"] ).strftime("%d/%m/%Y") - return accounts def delete_user_account(base_url: str, auth: str, stream: str, username: str) -> bool: - """Deletes a user account via the API. - - Args: - base_url: The base URL of the API. - auth: The authorization token. - stream: The stream associated with the account. - username: The username of the account to delete. - - Returns: - True if the account was deleted successfully, False otherwise. - """ - url = f"{base_url}/deleteAccount" + """Deletes a user account via the API.""" payload = {"stream": stream, "user": username} - headers = {"Authorization": f"Basic {auth}"} - response = requests.post(url, headers=headers, data=payload) - response.raise_for_status() - return "Deleted" in response.text + response = _make_api_request( + "POST", base_url, auth, "deleteAccount", payload=payload + ) + return response and "Deleted" in response.text -def add_user_account(base_url: str, auth: str, username: str, password: str, stream: str) -> bool: - """Adds a user account via the API. - - Args: - base_url: The base URL of the API. - auth: The authorization token. - username: The username of the new account. - password: The password for the new account. - stream: The stream to associate with the new account. - - Returns: - True if the account was added successfully, False otherwise. - """ - url = f"{base_url}/addAccount" +def add_user_account( + base_url: str, auth: str, username: str, password: str, stream: str +) -> bool: + """Adds a user account via the API.""" payload = {"username": username, "password": password, "stream": stream} - headers = {"Authorization": f"Basic {auth}"} - response = requests.post(url, headers=headers, data=payload) - return response.status_code == 200 + response = _make_api_request( + "POST", base_url, auth, "addAccount", payload=payload + ) + return response and response.status_code == 200 def get_stream_names(base_url: str, auth: str) -> List[str]: - """Retrieves a list of stream names from the API. - - Args: - base_url: The base URL of the API. - auth: The authorization token. - - Returns: - A list of stream names, or an empty list if an error occurs. - """ - url = f"{base_url}/getStreamNames" - headers = {"Authorization": f"Basic {auth}"} - try: - response = requests.get(url, headers=headers) - response.raise_for_status() - return response.json() - except (requests.exceptions.RequestException, json.JSONDecodeError): - return [] + """Retrieves a list of stream names from the API.""" + response = _make_api_request("GET", base_url, auth, "getStreamNames") + return response.json() if response else []