27 lines
846 B
Python
27 lines
846 B
Python
|
import base64
|
||
|
import os
|
||
|
|
||
|
from mlflow.tracking.request_auth.abstract_request_auth_provider import (
|
||
|
RequestAuthProvider,
|
||
|
)
|
||
|
|
||
|
|
||
|
class ProxyAuthProvider(RequestAuthProvider):
|
||
|
def __init__(self):
|
||
|
self.username = os.getenv("MLFLOW_TRACKING_USERNAME")
|
||
|
self.password = os.getenv("MLFLOW_TRACKING_PASSWORD")
|
||
|
|
||
|
def get_name(self):
|
||
|
return "proxy_auth_provider"
|
||
|
|
||
|
def get_auth(self):
|
||
|
# Add your Authelia Proxy-Authorization logic here
|
||
|
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}
|