from oauthlib.oauth2 import LegacyApplicationClient

from requests_oauthlib import OAuth2Session

class ProsodyRestClient(LegacyApplicationClient):

    pass

class ProsodyRestSession(OAuth2Session):

    def __init__(self, base_url=None, token_url=None, rest_url=None, *args, **kwargs):

        if base_url and not token_url:

            token_url = base_url + "/oauth2/token"

        if base_url and not rest_url:

            rest_url = base_url + "/rest"

        self._prosody_rest_url = rest_url

        self._prosody_token_url = token_url

        super().__init__(client=ProsodyRestClient(*args, **kwargs))

    def fetch_token(self, *args, **kwargs):

        return super().fetch_token(token_url=self._prosody_token_url, *args, **kwargs)

    def xmpp(self, json=None, *args, **kwargs):

        return self.post(self._prosody_rest_url, json=json, *args, **kwargs)

if __name__ == "__main__":

    # Example usage

    # from prosody_oauth import ProsodyRestSession

    from getpass import getpass

    p = ProsodyRestSession(base_url=input("Base URL: "), client_id="app")

    p.fetch_token(username=input("XMPP Address: "), password=getpass("Password: "))

    print(p.xmpp(json={"disco": True, "to": "jabber.org"}).json())