diff --git a/cps/kobo.py b/cps/kobo.py index 5f9db85a..688c11f0 100644 --- a/cps/kobo.py +++ b/cps/kobo.py @@ -19,7 +19,6 @@ import base64 from datetime import datetime, timezone -from functools import wraps import os import uuid import zipfile @@ -28,21 +27,30 @@ import json from urllib.parse import unquote import requests -from flask import Blueprint, request, make_response, jsonify, current_app, url_for, redirect, abort, g +from flask import ( + Blueprint, + request, + make_response, + jsonify, + current_app, + url_for, + redirect, + abort +) +from .cw_login import current_user from werkzeug.datastructures import Headers from sqlalchemy import func from sqlalchemy.sql.expression import and_, or_ from sqlalchemy.exc import StatementError -from flask_limiter.util import get_remote_address -from .cw_login import current_user, login_user from . import config, logger, kobo_auth, db, calibre_db, helper, shelf as shelf_lib, ub, csrf, kobo_sync_status +from . import isoLanguages, limiter from .epub import get_epub_layout from .constants import COVER_THUMBNAIL_SMALL, COVER_THUMBNAIL_MEDIUM, COVER_THUMBNAIL_LARGE, BASE_DIR from .helper import get_download_link from .services import SyncToken as SyncToken from .web import download_required -from . import isoLanguages, limiter +from .kobo_auth import requires_kobo_auth, get_auth_token KOBO_FORMATS = {"KEPUB": ["KEPUB"], "EPUB": ["EPUB3", "EPUB"]} KOBO_STOREAPI_URL = "https://storeapi.kobo.com" @@ -57,33 +65,6 @@ kobo_auth.register_url_value_preprocessor(kobo) log = logger.create() -def get_auth_token(): - if "auth_token" in g: - return g.get("auth_token") - else: - return None - - -def requires_kobo_auth(f): - @wraps(f) - def inner(*args, **kwargs): - auth_token = get_auth_token() - if auth_token is not None: - user = ( - ub.session.query(ub.User) - .join(ub.RemoteAuthToken) - .filter(ub.RemoteAuthToken.auth_token == auth_token).filter(ub.RemoteAuthToken.token_type==1) - .first() - ) - if user is not None: - login_user(user) - [limiter.limiter.clear(limit.limit, *limit.request_args) for limit in limiter.current_limits] - return f(*args, **kwargs) - log.debug("Received Kobo request without a recognizable auth token.") - return abort(401) - return inner - - def get_store_url_for_current_request(): # Programmatically modify the current url to point to the official Kobo store __, __, request_path_with_auth_token = request.full_path.rpartition("/kobo/") @@ -1339,5 +1320,3 @@ def NATIVE_KOBO_RESOURCES(): "userguide_host": "https://ereaderfiles.kobo.com", "wishlist_page": "https://www.kobo.com/{region}/{language}/account/wishlist" } - -limiter.limit("3/minute", key_func=get_remote_address)(kobo) diff --git a/cps/kobo_auth.py b/cps/kobo_auth.py index 042fb1c2..ae9d1c12 100644 --- a/cps/kobo_auth.py +++ b/cps/kobo_auth.py @@ -62,12 +62,13 @@ particular calls to non-Kobo specific endpoints such as the CalibreWeb book down from binascii import hexlify from datetime import datetime from os import urandom +from functools import wraps -from flask import g, Blueprint, request -from .cw_login import current_user +from flask import g, Blueprint, abort, request +from .cw_login import login_user, current_user from flask_babel import gettext as _ -from . import logger, config, calibre_db, db, helper, ub, lm +from . import logger, config, calibre_db, db, helper, ub, lm, limiter from .render_template import render_title_template from .usermanagement import user_login_required @@ -133,6 +134,13 @@ def disable_failed_auth_redirect_for_blueprint(bp): lm.blueprint_login_views[bp.name] = None +def get_auth_token(): + if "auth_token" in g: + return g.get("auth_token") + else: + return None + + def register_url_value_preprocessor(kobo): @kobo.url_value_preprocessor # pylint: disable=unused-variable @@ -140,3 +148,21 @@ def register_url_value_preprocessor(kobo): g.auth_token = values.pop("auth_token") +def requires_kobo_auth(f): + @wraps(f) + def inner(*args, **kwargs): + auth_token = get_auth_token() + if auth_token is not None: + user = ( + ub.session.query(ub.User) + .join(ub.RemoteAuthToken) + .filter(ub.RemoteAuthToken.auth_token == auth_token).filter(ub.RemoteAuthToken.token_type==1) + .first() + ) + if user is not None: + login_user(user) + [limiter.limiter.clear(limit.limit, *limit.request_args) for limit in limiter.current_limits] + return f(*args, **kwargs) + log.debug("Received Kobo request without a recognizable auth token.") + return abort(401) + return inner diff --git a/cps/main.py b/cps/main.py index 4b0e5440..b0f56b6d 100644 --- a/cps/main.py +++ b/cps/main.py @@ -46,10 +46,11 @@ def main(): try: from .kobo import kobo, get_kobo_activated from .kobo_auth import kobo_auth + from flask_limiter.util import get_remote_address kobo_available = get_kobo_activated() except (ImportError, AttributeError): # Catch also error for not installed flask-WTF (missing csrf decorator) kobo_available = False - kobo = kobo_auth = None + kobo = kobo_auth = get_remote_address = None try: from .oauth_bb import oauth @@ -76,6 +77,7 @@ def main(): app.register_blueprint(gdrive) app.register_blueprint(editbook) if kobo_available: + limiter.limit("3/minute", key_func=get_remote_address)(kobo) app.register_blueprint(kobo) app.register_blueprint(kobo_auth) if oauth_available: