diff --git a/spaceauth/__init__.py b/spaceauth/__init__.py index 459e668..932109b 100644 --- a/spaceauth/__init__.py +++ b/spaceauth/__init__.py @@ -4,7 +4,7 @@ from flask_login import LoginManager, login_user, logout_user, current_user, log from spaceauth.caps import cap_required -class SpaceAuth(object): +class SpaceAuth(LoginManager): def __init__(self, app=None, *args, **kwargs): self.oauth = OAuth() self.remote = self.oauth.remote_app( @@ -17,22 +17,22 @@ class SpaceAuth(object): self.remote.tokengetter(self.tokengetter) bp = Blueprint('spaceauth', __name__) - bp.add_url_rule('/login', 'login', self.login_view) - bp.add_url_rule('/logout', 'logout', self.logout_view) - bp.add_url_rule('/callback', 'callback', self.callback_view) + bp.add_url_rule('/login', 'login', self.login_view_handler) + bp.add_url_rule('/logout', 'logout', self.logout_view_handler) + bp.add_url_rule('/callback', 'callback', self.callback_view_handler) self.blueprint = bp - self.login_manager = LoginManager() - self.login_manager.refresh_view = 'spaceauth.login' - self.login_manager.login_view = 'spaceauth.login' - self.login_manager.user_loader(self.user_loader_handler) + super(SpaceAuth, self).__init__() + self.refresh_view = 'spaceauth.login' + self.login_view = 'spaceauth.login' + self.user_loader(self.user_loader_handler) if app: self.init_app(app, *args, **kwargs) def init_app(self, app, url_prefix='/oauth'): self.oauth.init_app(app) - self.login_manager.init_app(app) + super(SpaceAuth, self).init_app(app) app.register_blueprint(self.blueprint, url_prefix=url_prefix) @app.errorhandler(OAuthException) @@ -40,20 +40,20 @@ class SpaceAuth(object): flash('OAuth error occured', 'error') return redirect('/') - def login_view(self): + def login_view_handler(self): session['spaceauth_next'] = request.args.get('next') or request.referrer return self.remote.authorize( callback=url_for('spaceauth.callback', _external=True) ) - def logout_view(self): + def logout_view_handler(self): # TODO revoke token session.pop('spaceauth_token', None) session.pop('spaceauth_next', None) logout_user() return redirect('/') - def callback_view(self): + def callback_view_handler(self): resp = self.remote.authorized_response() if resp is None: raise OAuthException( @@ -79,16 +79,5 @@ class SpaceAuth(object): user.id = uid return user - def user_loader(self, func): - """ - Define flask_login-like user loader. Application is supposed to create - its user model when missing. Additional `profile` argument is passed - with user profile information right after login. - """ - - self.user_loader_handler = func - self.login_manager.user_loader(self.user_loader_handler) - return func - def user_profile(self): return self.remote.get('profile').data