From 77c318b32a58602a134247e2ede378524daef24c Mon Sep 17 00:00:00 2001
From: LordMathis <matus@namesny.com>
Date: Wed, 13 Dec 2023 23:25:44 +0100
Subject: [PATCH] Return Auth object instead of dict

---
 .../proxy_auth_header_provider.py               | 17 ++++++++++++-----
 setup.py                                        |  2 +-
 2 files changed, 13 insertions(+), 6 deletions(-)

diff --git a/mlflow_plugin_proxy_auth/proxy_auth_header_provider.py b/mlflow_plugin_proxy_auth/proxy_auth_header_provider.py
index 4b595fb..4f983c4 100644
--- a/mlflow_plugin_proxy_auth/proxy_auth_header_provider.py
+++ b/mlflow_plugin_proxy_auth/proxy_auth_header_provider.py
@@ -4,6 +4,7 @@ import os
 from mlflow.tracking.request_auth.abstract_request_auth_provider import (
     RequestAuthProvider,
 )
+from requests.auth import AuthBase
 
 
 class ProxyAuthProvider(RequestAuthProvider):
@@ -15,12 +16,18 @@ class ProxyAuthProvider(RequestAuthProvider):
         return "proxy_auth_provider"
 
     def get_auth(self):
-        # Add your Authelia Proxy-Authorization logic here
+        return ProxyAuth(self.username, self.password)
+
+
+class ProxyAuth(AuthBase):
+    def __init__(self, username, password):
+        self.username = username
+        self.password = password
+
+    def __call__(self, r):
         credentials = f"{self.username}:{self.password}"
         encoded_credentials = base64.b64encode(credentials.encode("utf-8")).decode(
             "utf-8"
         )
-        proxy_authorization_header = f"Basic {encoded_credentials}"
-
-        # Return a dictionary containing the header for authentication
-        return {"Proxy-Authorization": proxy_authorization_header}
+        r.headers["Proxy-Authorization"] = f"Basic {encoded_credentials}"
+        return r
diff --git a/setup.py b/setup.py
index 96ef14c..9334080 100644
--- a/setup.py
+++ b/setup.py
@@ -2,7 +2,7 @@ from setuptools import find_packages, setup
 
 setup(
     name="mlflow-plugin-proxy-auth",
-    version="0.0.2",
+    version="0.0.3",
     packages=find_packages(),
     install_requires=["mlflow"],
     entry_points={