From 348970946bc38df93d4aea2f7add9da176f51057 Mon Sep 17 00:00:00 2001 From: Santiago Lamora Date: Tue, 5 Apr 2022 10:36:04 +0200 Subject: [PATCH] Fix login_manager.user_loader Also redirect unathorized requests to login view --- ereuse_devicehub/devicehub.py | 170 ++++++++++++++++++++++------------ 1 file changed, 110 insertions(+), 60 deletions(-) diff --git a/ereuse_devicehub/devicehub.py b/ereuse_devicehub/devicehub.py index 7163173d..520ca132 100644 --- a/ereuse_devicehub/devicehub.py +++ b/ereuse_devicehub/devicehub.py @@ -10,7 +10,7 @@ from ereuse_utils.session import DevicehubClient from flask import _app_ctx_stack, g from flask_login import LoginManager, current_user from flask_sqlalchemy import SQLAlchemy -from teal.db import SchemaSQLAlchemy +from teal.db import ResourceNotFound, SchemaSQLAlchemy from teal.teal import Teal from ereuse_devicehub.auth import Auth @@ -29,32 +29,49 @@ class Devicehub(Teal): Dummy = Dummy jinja_environment = Environment - def __init__(self, - inventory: str, - config: DevicehubConfig = DevicehubConfig(), - db: SQLAlchemy = db, - import_name=__name__.split('.')[0], - static_url_path=None, - static_folder='static', - static_host=None, - host_matching=False, - subdomain_matching=False, - template_folder='templates', - instance_path=None, - instance_relative_config=False, - root_path=None, - Auth: Type[Auth] = Auth): + def __init__( + self, + inventory: str, + config: DevicehubConfig = DevicehubConfig(), + db: SQLAlchemy = db, + import_name=__name__.split('.')[0], + static_url_path=None, + static_folder='static', + static_host=None, + host_matching=False, + subdomain_matching=False, + template_folder='templates', + instance_path=None, + instance_relative_config=False, + root_path=None, + Auth: Type[Auth] = Auth, + ): assert inventory - super().__init__(config, db, inventory, import_name, static_url_path, static_folder, - static_host, - host_matching, subdomain_matching, template_folder, instance_path, - instance_relative_config, root_path, False, Auth) + super().__init__( + config, + db, + inventory, + import_name, + static_url_path, + static_folder, + static_host, + host_matching, + subdomain_matching, + template_folder, + instance_path, + instance_relative_config, + root_path, + False, + Auth, + ) self.id = inventory """The Inventory ID of this instance. In Teal is the app.schema.""" self.dummy = Dummy(self) - @self.cli.group(short_help='Inventory management.', - help='Manages the inventory {}.'.format(os.environ.get('dhi'))) + @self.cli.group( + short_help='Inventory management.', + help='Manages the inventory {}.'.format(os.environ.get('dhi')), + ) def inv(): pass @@ -69,43 +86,68 @@ class Devicehub(Teal): # configure Flask-Login login_manager = LoginManager() login_manager.init_app(self) + login_manager.login_view = "core.login" @login_manager.user_loader def load_user(user_id): - return User.query.get(user_id) + # TODO(@slamora) refactor when teal library has been drop. + # `load_user` expects None if the user ID is invalid or the + # session has expired so we need to handle Exception raised + # by teal (it's overriding default behaviour of flask-sqlalchemy + # which already returns None) + try: + return User.query.get(user_id) + except ResourceNotFound: + return None # noinspection PyMethodOverriding - @click.option('--name', '-n', - default='Test 1', - help='The human name of the inventory.') - @click.option('--org-name', '-on', - default='My Organization', - help='The name of the default organization that owns this inventory.') - @click.option('--org-id', '-oi', - default='foo-bar', - help='The Tax ID of the organization.') - @click.option('--tag-url', '-tu', - type=ereuse_utils.cli.URL(scheme=True, host=True, path=False), - default='http://example.com', - help='The base url (scheme and host) of the tag provider.') - @click.option('--tag-token', '-tt', - type=click.UUID, - default='899c794e-1737-4cea-9232-fdc507ab7106', - help='The token provided by the tag provider. It is an UUID.') - @click.option('--erase/--no-erase', - default=False, - help='Delete the schema before? ' - 'If --common is set this includes the common database.') - @click.option('--common/--no-common', - default=False, - help='Creates common databases. Only execute if the database is empty.') - def init_db(self, name: str, - org_name: str, - org_id: str, - tag_url: boltons.urlutils.URL, - tag_token: uuid.UUID, - erase: bool, - common: bool): + @click.option( + '--name', '-n', default='Test 1', help='The human name of the inventory.' + ) + @click.option( + '--org-name', + '-on', + default='My Organization', + help='The name of the default organization that owns this inventory.', + ) + @click.option( + '--org-id', '-oi', default='foo-bar', help='The Tax ID of the organization.' + ) + @click.option( + '--tag-url', + '-tu', + type=ereuse_utils.cli.URL(scheme=True, host=True, path=False), + default='http://example.com', + help='The base url (scheme and host) of the tag provider.', + ) + @click.option( + '--tag-token', + '-tt', + type=click.UUID, + default='899c794e-1737-4cea-9232-fdc507ab7106', + help='The token provided by the tag provider. It is an UUID.', + ) + @click.option( + '--erase/--no-erase', + default=False, + help='Delete the schema before? ' + 'If --common is set this includes the common database.', + ) + @click.option( + '--common/--no-common', + default=False, + help='Creates common databases. Only execute if the database is empty.', + ) + def init_db( + self, + name: str, + org_name: str, + org_id: str, + tag_url: boltons.urlutils.URL, + tag_token: uuid.UUID, + erase: bool, + common: bool, + ): """Creates an inventory. This creates the database and adds the inventory to the @@ -120,10 +162,14 @@ class Devicehub(Teal): with click_spinner.spinner(): if erase: self.db.drop_all(common_schema=common) - assert not db.has_schema(self.id), 'Schema {} already exists.'.format(self.id) + assert not db.has_schema(self.id), 'Schema {} already exists.'.format( + self.id + ) exclude_schema = 'common' if not common else None self._init_db(exclude_schema=exclude_schema) - InventoryDef.set_inventory_config(name, org_name, org_id, tag_url, tag_token) + InventoryDef.set_inventory_config( + name, org_name, org_id, tag_url, tag_token + ) DeviceSearch.set_all_devices_tokens_if_empty(self.db.session) self._init_resources(exclude_schema=exclude_schema) self.db.session.commit() @@ -138,8 +184,11 @@ class Devicehub(Teal): return True - @click.confirmation_option(prompt='Are you sure you want to delete the inventory {}?' - .format(os.environ.get('dhi'))) + @click.confirmation_option( + prompt='Are you sure you want to delete the inventory {}?'.format( + os.environ.get('dhi') + ) + ) def delete_inventory(self): """Erases an inventory. @@ -161,8 +210,9 @@ class Devicehub(Teal): def _prepare_request(self): """Prepares request stuff.""" inv = g.inventory = Inventory.current # type: Inventory - g.tag_provider = DevicehubClient(base_url=inv.tag_provider, - token=DevicehubClient.encode_token(inv.tag_token)) + g.tag_provider = DevicehubClient( + base_url=inv.tag_provider, token=DevicehubClient.encode_token(inv.tag_token) + ) # NOTE: models init methods expects that current user is # available on g.user (e.g. to initialize object owner) g.user = current_user