From 97e0c353e11df88d218f3b934aab32adb8c65568 Mon Sep 17 00:00:00 2001 From: Gabryel Reyes Date: Fri, 21 Jun 2024 11:47:48 +0200 Subject: [PATCH 1/3] Created Superset request wrapper --- src/pySupersetCli/__main__.py | 31 ++++- src/pySupersetCli/superset.py | 212 ++++++++++++++++++++++++++++++++++ 2 files changed, 241 insertions(+), 2 deletions(-) create mode 100644 src/pySupersetCli/superset.py diff --git a/src/pySupersetCli/__main__.py b/src/pySupersetCli/__main__.py index d01e1c1..da1a305 100644 --- a/src/pySupersetCli/__main__.py +++ b/src/pySupersetCli/__main__.py @@ -39,6 +39,7 @@ from pySupersetCli.version import __version__, __author__, __email__, __repository__, __license__ from pySupersetCli.ret import Ret +from pySupersetCli.superset import Superset ################################################################################ @@ -115,6 +116,14 @@ def add_parser() -> argparse.ArgumentParser: help="Print full command details before executing the command.\ Enables logs of type INFO and WARNING.") + parser.add_argument("--no_ssl", + action="store_true", + help="Disables SSL certificate verification.") + + parser.add_argument("--basic_auth", + action="store_true", + help="Use basic authentication instead of LDAP.") + return parser @@ -151,7 +160,25 @@ def main() -> Ret: for arg in vars(args): LOG.info("* %s = %s", arg, vars(args)[arg]) - if Ret.OK == ret_status: + # Create Superset client. + try: + verify_ssl = not args.no_ssl + provider = Superset.Provider.LDAP + + if args.basic_auth: + provider = Superset.Provider.DB + + # pylint: disable=unused-variable + client = Superset(args.server, + args.user, + args.password, + provider, + verify_ssl=verify_ssl) + + except RuntimeError as e: + LOG.error("Failed to create Superset client: %s", e) + ret_status = Ret.ERROR_LOGIN + else: handler = None # Find the command handler. @@ -162,7 +189,7 @@ def main() -> Ret: # Execute the command. if handler is not None: - ret_status = Ret.OK + ret_status = handler(args, client) else: LOG.error("Command '%s' not found!", args.cmd) ret_status = Ret.ERROR_INVALID_ARGUMENTS diff --git a/src/pySupersetCli/superset.py b/src/pySupersetCli/superset.py new file mode 100644 index 0000000..0c24a95 --- /dev/null +++ b/src/pySupersetCli/superset.py @@ -0,0 +1,212 @@ +"""Server wrapper for requests to the Superset API.""" + +# BSD 3-Clause License +# +# Copyright (c) 2024, NewTec GmbH +# +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: +# +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. +# +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. +# +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. +# +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICU5LAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +################################################################################ +# Imports +################################################################################ + +from dataclasses import dataclass +import logging +import requests +import urllib3 + + +################################################################################ +# Variables +################################################################################ + +LOG: logging.Logger = logging.getLogger(__name__) + +################################################################################ +# Classes +################################################################################ + + +class Superset: # pylint: disable=too-few-public-methods + """ + Wrapper of the requests module for the Superset API. + Handles the authentication and the API calls. + Implements parts of the Superset API: https://superset.apache.org/docs/api/ + """ + + @dataclass + class Provider: + """ + Enum for the supported authentication providers. + """ + DB = "db" + LDAP = "ldap" + + # pylint: disable=too-many-arguments + def __init__(self, + server_url: str, + username: str, + password: str, + provider: Provider, + verify_ssl: bool = True) -> None: + """ + Initializes the Superset object and logs in the user. + + Args: + server_url (str): The URL of the Superset server. + username (str): The username of the user. + password (str): The password of the user. + provider (Provider): The authentication provider. + verify_ssl (bool): Verify the SSL certificate of the server. + """ + self._server_url: str = f"{server_url}/api/v1" + self._access_token: str = "" + self._csrf_token: str = "" + self._timeout: int = 60 + self._verify_ssl: bool = verify_ssl + + if not self._verify_ssl: + # Disable SSL warnings if SSL verification is disabled + urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning) + + # Login the user and retrieve the access token and the CSRF token + self._login(username, password, provider) + + def _login(self, username: str, password: str, provider: Provider) -> None: + """ + Logs in the user and retrieves the access token and the refresh token. + + Args: + username (str): The username of the user. + password (str): The password of the user. + provider (Provider): The authentication provider. + + Returns: + None + """ + login_endpoint: str = "/security/login" + crsf_token_endpoint: str = "/security/csrf_token/" + + login_body: dict = { + "password": password, + "provider": provider, + "refresh": True, + "username": username + } + + # Send the login request + ret_code, response = self.request("POST", + login_endpoint, + json=login_body) + + if ret_code != 200: + LOG.fatal("Login failed: %s", response.get("message")) + raise RuntimeError("Login failed") + + self._access_token = response.get("access_token", "") + + # Get the CSRF token + ret_code, response = self.request("GET", crsf_token_endpoint) + + if ret_code != 200: + LOG.fatal("Get CSRF token failed: %s", response.get("message")) + raise RuntimeError("Get CSRF token failed") + + self._csrf_token = response.get("result", "") + + if self._access_token == "" or self._csrf_token == "": + LOG.fatal("Tokens failed: Access token or CSRF token not received.") + raise RuntimeError("Tokens failed") + + def request(self, + method: str, + endpoint: str, + **request_kwargs) -> tuple[int, dict]: + """ + Sends a request to the Superset API. + + Args: + method (str): The HTTP method of the request. + endpoint (str): The endpoint of the request after '/api/v1'. + data (dict): The data of the request. + request_kwargs (dict): Additional keyword arguments for the request. + Can be any accepted by the Requests module. + + Returns: + dict: The response of the request. + """ + + url: str = f"{self._server_url}{endpoint}" + headers: dict = {} + response_code: int = 0 + reponse_data: dict = {} + + # If already logged in, add the access token to the headers + if self._access_token != "": + headers = { + 'Authorization': f'Bearer {self._access_token}', + 'referer': self._server_url, + 'X-CSRFToken': self._csrf_token + } + + try: + # Send the request + response: requests.Response = requests.request( + method=method, + url=url, + headers=headers, + timeout=self._timeout, + verify=self._verify_ssl, + ** request_kwargs) + except requests.exceptions.SSLError as e: + LOG.error("SSL error: %s", e) + + if self._verify_ssl is True: + LOG.error("If you trust the server you are connecting to (%s), " + + "consider deactivating SSL verification.", self._server_url) + + else: + # Check if the token has expired + if (reponse_data.get('message') == "Token has expired") and \ + (response_code == 401): + LOG.error("Refreshing token is not implemented.") + else: + response_code = response.status_code + reponse_data = response.json() + + LOG.info("Request: %s %s", method, url) + LOG.info("Response Code: %s", response_code) + + return (response_code, reponse_data) + + +################################################################################ +# Functions +################################################################################ + +################################################################################ +# Main +################################################################################ From 2a6a3175e9463b3fefa171f5df9afe02696f70a0 Mon Sep 17 00:00:00 2001 From: Gabryel Reyes Date: Fri, 21 Jun 2024 13:24:32 +0200 Subject: [PATCH 2/3] Removed magic numbers --- src/pySupersetCli/superset.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/pySupersetCli/superset.py b/src/pySupersetCli/superset.py index 0c24a95..ad2ba97 100644 --- a/src/pySupersetCli/superset.py +++ b/src/pySupersetCli/superset.py @@ -122,7 +122,7 @@ def _login(self, username: str, password: str, provider: Provider) -> None: login_endpoint, json=login_body) - if ret_code != 200: + if requests.codes.ok != ret_code: LOG.fatal("Login failed: %s", response.get("message")) raise RuntimeError("Login failed") @@ -131,7 +131,7 @@ def _login(self, username: str, password: str, provider: Provider) -> None: # Get the CSRF token ret_code, response = self.request("GET", crsf_token_endpoint) - if ret_code != 200: + if requests.codes.ok != ret_code: LOG.fatal("Get CSRF token failed: %s", response.get("message")) raise RuntimeError("Get CSRF token failed") @@ -191,7 +191,7 @@ def request(self, else: # Check if the token has expired if (reponse_data.get('message') == "Token has expired") and \ - (response_code == 401): + (requests.codes.unauthorized == response_code): LOG.error("Refreshing token is not implemented.") else: response_code = response.status_code From 528f5ee6cfaee61c0b8e77dcf3ab400f75cac115 Mon Sep 17 00:00:00 2001 From: Gabryel Reyes Date: Fri, 21 Jun 2024 13:32:00 +0200 Subject: [PATCH 3/3] Disabled warnings --- src/pySupersetCli/superset.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/src/pySupersetCli/superset.py b/src/pySupersetCli/superset.py index ad2ba97..3acc52a 100644 --- a/src/pySupersetCli/superset.py +++ b/src/pySupersetCli/superset.py @@ -122,7 +122,7 @@ def _login(self, username: str, password: str, provider: Provider) -> None: login_endpoint, json=login_body) - if requests.codes.ok != ret_code: + if requests.codes.ok != ret_code: # pylint: disable=no-member LOG.fatal("Login failed: %s", response.get("message")) raise RuntimeError("Login failed") @@ -131,7 +131,7 @@ def _login(self, username: str, password: str, provider: Provider) -> None: # Get the CSRF token ret_code, response = self.request("GET", crsf_token_endpoint) - if requests.codes.ok != ret_code: + if requests.codes.ok != ret_code: # pylint: disable=no-member LOG.fatal("Get CSRF token failed: %s", response.get("message")) raise RuntimeError("Get CSRF token failed") @@ -190,8 +190,8 @@ def request(self, else: # Check if the token has expired - if (reponse_data.get('message') == "Token has expired") and \ - (requests.codes.unauthorized == response_code): + if (requests.codes.unauthorized == response_code) and \ + (reponse_data.get('message') == "Token has expired"): # pylint: disable=no-member LOG.error("Refreshing token is not implemented.") else: response_code = response.status_code