calibre-web/cps/oauth.py

# -*- coding: utf-8 -*-

#  This file is part of the Calibre-Web (https://github.com/janeczku/calibre-web)
#    Copyright (C) 2018-2019 jim3ma
#
#  This program is free software: you can redistribute it and/or modify
#  it under the terms of the GNU General Public License as published by
#  the Free Software Foundation, either version 3 of the License, or
#  (at your option) any later version.
#
#  This program is distributed in the hope that it will be useful,
#  but WITHOUT ANY WARRANTY; without even the implied warranty of
#  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#  GNU General Public License for more details.
#
#  You should have received a copy of the GNU General Public License
#  along with this program. If not, see <http://www.gnu.org/licenses/>

from __future__ import division, print_function, unicode_literals
from flask import session


try:
    from flask_dance.consumer.backend.sqla import SQLAlchemyBackend, first, _get_real_user
    from sqlalchemy.orm.exc import NoResultFound
except ImportError:
    # fails on flask-dance >1.3, due to renaming
    try:
        from flask_dance.consumer.storage.sqla import SQLAlchemyStorage as SQLAlchemyBackend
        from flask_dance.consumer.storage.sqla import first, _get_real_user
        from sqlalchemy.orm.exc import NoResultFound
    except ImportError:
        pass

try:
    class OAuthBackend(SQLAlchemyBackend):
        """
        Stores and retrieves OAuth tokens using a relational database through
        the `SQLAlchemy`_ ORM.

        .. _SQLAlchemy: http://www.sqlalchemy.org/
        """
        def __init__(self, model, session, provider_id,
                     user=None, user_id=None, user_required=None, anon_user=None,
                     cache=None):
            self.provider_id = provider_id
            super(OAuthBackend, self).__init__(model, session, user, user_id, user_required, anon_user, cache)

        def get(self, blueprint, user=None, user_id=None):
            if self.provider_id + '_oauth_token' in session and session[self.provider_id + '_oauth_token'] != '':
                return session[blueprint.name + '_oauth_token']
            # check cache
            cache_key = self.make_cache_key(blueprint=blueprint, user=user, user_id=user_id)
            token = self.cache.get(cache_key)
            if token:
                return token

            # if not cached, make database queries
            query = (
                self.session.query(self.model)
                .filter_by(provider=self.provider_id)
            )
            uid = first([user_id, self.user_id, blueprint.config.get("user_id")])
            u = first(_get_real_user(ref, self.anon_user)
                      for ref in (user, self.user, blueprint.config.get("user")))

            use_provider_user_id = False
            if self.provider_id + '_oauth_user_id' in session and session[self.provider_id + '_oauth_user_id'] != '':
                query = query.filter_by(provider_user_id=session[self.provider_id + '_oauth_user_id'])
                use_provider_user_id = True

            if self.user_required and not u and not uid and not use_provider_user_id:
                #raise ValueError("Cannot get OAuth token without an associated user")
                return None
            # check for user ID
            if hasattr(self.model, "user_id") and uid:
                query = query.filter_by(user_id=uid)
            # check for user (relationship property)
            elif hasattr(self.model, "user") and u:
                query = query.filter_by(user=u)
            # if we have the property, but not value, filter by None
            elif hasattr(self.model, "user_id"):
                query = query.filter_by(user_id=None)
            # run query
            try:
                token = query.one().token
            except NoResultFound:
                token = None

            # cache the result
            self.cache.set(cache_key, token)

            return token

        def set(self, blueprint, token, user=None, user_id=None):
            uid = first([user_id, self.user_id, blueprint.config.get("user_id")])
            u = first(_get_real_user(ref, self.anon_user)
                          for ref in (user, self.user, blueprint.config.get("user")))

            if self.user_required and not u and not uid:
                raise ValueError("Cannot set OAuth token without an associated user")

            # if there was an existing model, delete it
            existing_query = (
                self.session.query(self.model)
                .filter_by(provider=self.provider_id)
            )
            # check for user ID
            has_user_id = hasattr(self.model, "user_id")
            if has_user_id and uid:
                existing_query = existing_query.filter_by(user_id=uid)
            # check for user (relationship property)
            has_user = hasattr(self.model, "user")
            if has_user and u:
                existing_query = existing_query.filter_by(user=u)
            # queue up delete query -- won't be run until commit()
            existing_query.delete()
            # create a new model for this token
            kwargs = {
                "provider": self.provider_id,
                "token": token,
            }
            if has_user_id and uid:
                kwargs["user_id"] = uid
            if has_user and u:
                kwargs["user"] = u
            self.session.add(self.model(**kwargs))
            # commit to delete and add simultaneously
            self.session.commit()
            # invalidate cache
            self.cache.delete(self.make_cache_key(
                blueprint=blueprint, user=user, user_id=user_id
            ))

        def delete(self, blueprint, user=None, user_id=None):
            query = (
                self.session.query(self.model)
                .filter_by(provider=self.provider_id)
            )
            uid = first([user_id, self.user_id, blueprint.config.get("user_id")])
            u = first(_get_real_user(ref, self.anon_user)
                      for ref in (user, self.user, blueprint.config.get("user")))

            if self.user_required and not u and not uid:
                raise ValueError("Cannot delete OAuth token without an associated user")

            # check for user ID
            if hasattr(self.model, "user_id") and uid:
                query = query.filter_by(user_id=uid)
            # check for user (relationship property)
            elif hasattr(self.model, "user") and u:
                query = query.filter_by(user=u)
            # if we have the property, but not value, filter by None
            elif hasattr(self.model, "user_id"):
                query = query.filter_by(user_id=None)
            # run query
            query.delete()
            self.session.commit()
            # invalidate cache
            self.cache.delete(self.make_cache_key(
                blueprint=blueprint, user=user, user_id=user_id,
            ))

except Exception:
    pass