from django_roa.db.query import RemoteQuerySet
from django_roa.db import get_roa_headers

from dockerhub.utils import get_jwt
from dockerhub.exceptions import AuthenticationFailure


class HubRemoteQuerySet(RemoteQuerySet):
    def __init__(self, username, password=None, model=None, query=None):
        self.username = username
        self.password = password
        super(HubRemoteQuerySet, self).__init__(model=model, query=query)

    def _get_http_headers(self):
        headers = get_roa_headers()
        try:
            headers.update(**{"Authorization": "JWT ".format(get_jwt(self.username, self.password))})
        except AuthenticationFailure:
            pass