Initial implementation of proxy auth header

This commit is contained in:
LordMathis 2023-12-13 22:03:07 +01:00
parent 6feec69db7
commit 9f2a5cdfbb
3 changed files with 44 additions and 0 deletions

7
.flake8 Normal file
View File

@ -0,0 +1,7 @@
[flake8]
exclude = .git,__pycache__,data,old,build,dist
max-complexity = 10
max-line-length = 88
extend-ignore = E203
per-file-ignores =
*/__init__.py: F401

View File

@ -0,0 +1,26 @@
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}

11
setup.py Normal file
View File

@ -0,0 +1,11 @@
from setuptools import find_packages, setup
setup(
name="mlflow-plugin-proxy-auth",
version="0.0.1",
packages=find_packages(),
install_requires=["mlflow"],
entry_points={
"mlflow.request_auth_provider": "dummy-backend=mlflow_plugin_proxy_auth.proxy_auth_header_provider:ProxyAuthProvider",
},
)