Use API authentication instead of token scraping

src/preset_cli/auth/main.py

@@ -2,12 +2,15 @@
 Mechanisms for authentication and authorization.
 """
 
+import logging
 from typing import Any, Dict
 
 from requests import Response, Session
 from requests.adapters import HTTPAdapter
 from urllib3.util import Retry
 
+_logger = logging.getLogger(__name__)
+
 
 class Auth:  # pylint: disable=too-few-public-methods
     """
@@ -46,6 +49,8 @@ def reauth(self, r: Response, *args: Any, **kwargs: Any) -> Response:
         if r.status_code != 401:
             return r
 
+        _logger.debug("Token expired. Re-authenticating...")
+
         try:
             self.auth()
         except NotImplementedError:

src/preset_cli/auth/superset.py

@@ -2,35 +2,99 @@
 Mechanisms for authentication and authorization for Superset instances.
 """
 
+import logging
 from typing import Dict, Optional
 
 from bs4 import BeautifulSoup
+from requests.exceptions import RequestException
 from yarl import URL
 
 from preset_cli.auth.main import Auth
 from preset_cli.auth.token import TokenAuth
 
+_logger = logging.getLogger(__name__)
+
 
 class UsernamePasswordAuth(Auth):  # pylint: disable=too-few-public-methods
     """
     Auth to Superset via username/password.
     """
 
-    def __init__(self, baseurl: URL, username: str, password: Optional[str] = None):
+    def __init__(
+        self,
+        baseurl: URL,
+        username: str,
+        password: Optional[str] = None,
+        provider: Optional[str] = "db",
+    ):
         super().__init__()
 
         self.csrf_token: Optional[str] = None
         self.baseurl = baseurl
         self.username = username
         self.password = password
+        self.provider = provider or "db"
         self.auth()
 
     def get_headers(self) -> Dict[str, str]:
         return {"X-CSRFToken": self.csrf_token} if self.csrf_token else {}
 
+    def get_access_token(self):
+        """
+        Get an access token from superset API: api/v1/security/login.
+        """
+        body = {
+            "username": self.username,
+            "password": self.password,
+            "provider": self.provider,
+        }
+        if "Referer" in self.session.headers:
+            del self.session.headers["Referer"]
+        response = self.session.post(self.baseurl / "api/v1/security/login", json=body)
+        response.raise_for_status()
+        return response.json()["access_token"]
+
+    def get_csrf_token(self):
+        """
+        Get a CSRF token from superset API: api/v1/security/csrf_token .
+        """
+        response = self.session.get(self.baseurl / "api/v1/security/csrf_token/")
+        response.raise_for_status()
+        return response.json()["result"]
+
     def auth(self) -> None:
         """
         Login to get CSRF token and cookies.
+
+        Try the documented REST API first; fall back to the legacy HTML-scraping
+        flow if the API is unavailable (e.g. on older Superset versions that do
+        not expose ``/api/v1/security/login``).
+        """
+        try:
+            self.session.headers["Authorization"] = f"Bearer {self.get_access_token()}"
+            csrf_token = self.get_csrf_token()
+        except (RequestException, KeyError, ValueError) as ex:
+            _logger.warning(
+                "API authentication failed (%s); falling back to legacy "
+                "HTML-based login flow.",
+                ex,
+            )
+            self.session.headers.pop("Authorization", None)
+            self._legacy_auth()
+            return
+
+        if csrf_token:
+            self.session.headers["X-CSRFToken"] = csrf_token
+            self.session.headers["Referer"] = str(
+                self.baseurl / "api/v1/security/csrf_token/",
+            )
+            self.csrf_token = csrf_token
+
+    def _legacy_auth(self) -> None:
+        """
+        Legacy login flow: scrape the CSRF token from ``/login/`` and POST
+        credentials as form data. Kept as a fallback for older Superset
+        instances that don't expose the security API.
         """
         data = {"username": self.username, "password": self.password}
 
@@ -43,7 +107,6 @@ def auth(self) -> None:
             data["csrf_token"] = csrf_token
             self.csrf_token = csrf_token
 
-        # set cookies
         self.session.post(self.baseurl / "login/", data=data)
 
 

src/preset_cli/cli/superset/main.py

@@ -37,6 +37,11 @@
     help="Password (leave empty for prompt)",
 )
 @click.option("--loglevel", default="INFO")
+@click.option(
+    "--provider",
+    default="db",
+    help="Provider that store the superset credentials (e.g. ldap, db, etc.)",
+)
 @click.version_option()
 @click.pass_context
 def superset_cli(  # pylint: disable=too-many-arguments
@@ -46,6 +51,7 @@ def superset_cli(  # pylint: disable=too-many-arguments
     username: str,
     password: str,
     loglevel: str,
+    provider: str,
 ):
     """
     An Apache Superset CLI.
@@ -61,7 +67,12 @@ def superset_cli(  # pylint: disable=too-many-arguments
         if jwt_token:
             ctx.obj["AUTH"] = SupersetJWTAuth(jwt_token, URL(instance))
         else:
-            ctx.obj["AUTH"] = UsernamePasswordAuth(URL(instance), username, password)
+            ctx.obj["AUTH"] = UsernamePasswordAuth(
+                URL(instance),
+                username,
+                password,
+                provider,
+            )
 
 
 superset_cli.add_command(sql)

tests/auth/superset_test.py

@@ -14,11 +14,15 @@ def test_username_password_auth(requests_mock: Mocker) -> None:
     Tests for the username/password authentication mechanism.
     """
     csrf_token = "CSFR_TOKEN"
+    access_token = "ACCESS_TOKEN"
     requests_mock.get(
-        "https://superset.example.org/login/",
-        text=f'<html><body><input id="csrf_token" value="{csrf_token}"></body></html>',
+        "https://superset.example.org/api/v1/security/csrf_token/",
+        json={"result": csrf_token},
+    )
+    requests_mock.post(
+        "https://superset.example.org/api/v1/security/login",
+        json={"access_token": access_token},
     )
-    requests_mock.post("https://superset.example.org/login/")
 
     auth = UsernamePasswordAuth(
         URL("https://superset.example.org/"),
@@ -29,21 +33,20 @@ def test_username_password_auth(requests_mock: Mocker) -> None:
         "X-CSRFToken": csrf_token,
     }
 
-    assert (
-        requests_mock.last_request.text
-        == "username=admin&password=password123&csrf_token=CSFR_TOKEN"
-    )
-
 
 def test_username_password_auth_no_csrf(requests_mock: Mocker) -> None:
     """
     Tests for the username/password authentication mechanism.
     """
+    access_token = "ACCESS_TOKEN"
     requests_mock.get(
-        "https://superset.example.org/login/",
-        text="<html><body>WTF_CSRF_ENABLED = False</body></html>",
+        "https://superset.example.org/api/v1/security/csrf_token/",
+        json={"result": None},
+    )
+    requests_mock.post(
+        "https://superset.example.org/api/v1/security/login",
+        json={"access_token": access_token},
     )
-    requests_mock.post("https://superset.example.org/login/")
 
     auth = UsernamePasswordAuth(
         URL("https://superset.example.org/"),
@@ -53,7 +56,33 @@ def test_username_password_auth_no_csrf(requests_mock: Mocker) -> None:
     # pylint: disable=use-implicit-booleaness-not-comparison
     assert auth.get_headers() == {}
 
-    assert requests_mock.last_request.text == "username=admin&password=password123"
+
+def test_username_password_auth_legacy_fallback(requests_mock: Mocker) -> None:
+    """
+    When the security API is unavailable, fall back to the legacy
+    HTML-scraping login flow.
+    """
+    csrf_token = "LEGACY_CSRF"
+    requests_mock.post(
+        "https://superset.example.org/api/v1/security/login",
+        status_code=404,
+    )
+    requests_mock.get(
+        "https://superset.example.org/login/",
+        text=(
+            f'<html><body><input id="csrf_token" name="csrf_token" '
+            f'value="{csrf_token}" /></body></html>'
+        ),
+    )
+    requests_mock.post("https://superset.example.org/login/")
+
+    auth = UsernamePasswordAuth(
+        URL("https://superset.example.org/"),
+        "admin",
+        "password123",
+    )
+    assert auth.get_headers() == {"X-CSRFToken": csrf_token}
+    assert "Authorization" not in auth.session.headers
 
 
 def test_jwt_auth_superset(mocker: MockerFixture) -> None:

tests/cli/superset/main_test.py

@@ -163,3 +163,49 @@ def test_superset_jwt_auth(mocker: MockerFixture) -> None:
     )
 
     SupersetJWTAuth.assert_called_with("SECRET", URL("http://localhost:8088/"))
+
+
+def test_superset_cli_default_provider(mocker: MockerFixture) -> None:
+    """
+    Test that the --provider option defaults to 'db'.
+    """
+    username_password_auth = mocker.patch(
+        "preset_cli.cli.superset.main.UsernamePasswordAuth",
+    )
+
+    runner = CliRunner()
+    runner.invoke(
+        superset_cli,
+        ["http://localhost:8088/", "export"],
+        catch_exceptions=False,
+    )
+
+    username_password_auth.assert_called_with(
+        URL("http://localhost:8088/"),
+        "admin",
+        "admin",
+        "db",
+    )
+
+
+def test_superset_cli_provider(mocker: MockerFixture) -> None:
+    """
+    Test that --provider db is passed through to UsernamePasswordAuth.
+    """
+    username_password_auth = mocker.patch(
+        "preset_cli.cli.superset.main.UsernamePasswordAuth",
+    )
+
+    runner = CliRunner()
+    runner.invoke(
+        superset_cli,
+        ["--provider=ldap", "http://localhost:8088/", "export"],
+        catch_exceptions=False,
+    )
+
+    username_password_auth.assert_called_with(
+        URL("http://localhost:8088/"),
+        "admin",
+        "admin",
+        "ldap",
+    )