| # -*- coding: utf-8 -*- |
| # |
| # Licensed to the Apache Software Foundation (ASF) under one |
| # or more contributor license agreements. See the NOTICE file |
| # distributed with this work for additional information |
| # regarding copyright ownership. The ASF licenses this file |
| # to you under the Apache License, Version 2.0 (the |
| # "License"); you may not use this file except in compliance |
| # with the License. You may obtain a copy of the License at |
| # |
| # http://www.apache.org/licenses/LICENSE-2.0 |
| # |
| # Unless required by applicable law or agreed to in writing, |
| # software distributed under the License is distributed on an |
| # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY |
| # KIND, either express or implied. See the License for the |
| # specific language governing permissions and limitations |
| # under the License. |
| # |
| |
| from flask import g |
| from flask_appbuilder.security.sqla import models as sqla_models |
| from flask_appbuilder.security.sqla.manager import SecurityManager |
| from sqlalchemy import or_, and_ |
| |
| from airflow import models |
| from airflow.exceptions import AirflowException |
| from airflow.utils.db import provide_session |
| from airflow.utils.log.logging_mixin import LoggingMixin |
| from airflow.www_rbac.app import appbuilder |
| from airflow.www_rbac.utils import CustomSQLAInterface |
| |
| |
| EXISTING_ROLES = { |
| 'Admin', |
| 'Viewer', |
| 'User', |
| 'Op', |
| 'Public', |
| } |
| |
| |
| class AirflowSecurityManager(SecurityManager, LoggingMixin): |
| ########################################################################### |
| # VIEW MENUS |
| ########################################################################### |
| VIEWER_VMS = { |
| 'Airflow', |
| 'DagModelView', |
| 'Browse', |
| 'DAG Runs', |
| 'DagRunModelView', |
| 'Task Instances', |
| 'TaskInstanceModelView', |
| 'SLA Misses', |
| 'SlaMissModelView', |
| 'Jobs', |
| 'JobModelView', |
| 'Logs', |
| 'LogModelView', |
| 'Docs', |
| 'Documentation', |
| 'Github', |
| 'About', |
| 'Version', |
| 'VersionView', |
| } |
| |
| USER_VMS = VIEWER_VMS |
| |
| OP_VMS = { |
| 'Admin', |
| 'Configurations', |
| 'ConfigurationView', |
| 'Connections', |
| 'ConnectionModelView', |
| 'Pools', |
| 'PoolModelView', |
| 'Variables', |
| 'VariableModelView', |
| 'XComs', |
| 'XComModelView', |
| } |
| |
| ########################################################################### |
| # PERMISSIONS |
| ########################################################################### |
| |
| VIEWER_PERMS = { |
| 'menu_access', |
| 'can_index', |
| 'can_list', |
| 'can_show', |
| 'can_chart', |
| 'can_dag_stats', |
| 'can_dag_details', |
| 'can_task_stats', |
| 'can_code', |
| 'can_log', |
| 'can_get_logs_with_metadata', |
| 'can_tries', |
| 'can_graph', |
| 'can_tree', |
| 'can_task', |
| 'can_task_instances', |
| 'can_xcom', |
| 'can_gantt', |
| 'can_landing_times', |
| 'can_last_dagruns', |
| 'can_duration', |
| 'can_blocked', |
| 'can_rendered', |
| 'can_pickle_info', |
| 'can_version', |
| } |
| |
| USER_PERMS = { |
| 'can_dagrun_clear', |
| 'can_run', |
| 'can_trigger', |
| 'can_add', |
| 'can_edit', |
| 'can_delete', |
| 'can_failed', |
| 'can_paused', |
| 'can_refresh', |
| 'can_success', |
| 'muldelete', |
| 'set_failed', |
| 'set_running', |
| 'set_success', |
| 'clear', |
| 'can_clear', |
| } |
| |
| OP_PERMS = { |
| 'can_conf', |
| 'can_varimport', |
| } |
| |
| # global view-menu for dag-level access |
| DAG_VMS = { |
| 'all_dags' |
| } |
| |
| WRITE_DAG_PERMS = { |
| 'can_dag_edit', |
| } |
| |
| READ_DAG_PERMS = { |
| 'can_dag_read', |
| } |
| |
| DAG_PERMS = WRITE_DAG_PERMS | READ_DAG_PERMS |
| |
| ########################################################################### |
| # DEFAULT ROLE CONFIGURATIONS |
| ########################################################################### |
| |
| ROLE_CONFIGS = [ |
| { |
| 'role': 'Viewer', |
| 'perms': VIEWER_PERMS | READ_DAG_PERMS, |
| 'vms': VIEWER_VMS | DAG_VMS |
| }, |
| { |
| 'role': 'User', |
| 'perms': VIEWER_PERMS | USER_PERMS | DAG_PERMS, |
| 'vms': VIEWER_VMS | DAG_VMS | USER_VMS, |
| }, |
| { |
| 'role': 'Op', |
| 'perms': VIEWER_PERMS | USER_PERMS | OP_PERMS | DAG_PERMS, |
| 'vms': VIEWER_VMS | DAG_VMS | USER_VMS | OP_VMS, |
| }, |
| ] |
| |
| def __init__(self, appbuilder): |
| super(AirflowSecurityManager, self).__init__(appbuilder) |
| |
| # Go and fix up the SQLAInterface used from the stock one to our subclass. |
| # This is needed to support the "hack" where we had to edit |
| # FieldConverter.conversion_table in place in airflow.www.utils |
| for attr in dir(self): |
| if not attr.endswith('view'): |
| continue |
| view = getattr(self, attr, None) |
| if not view or not getattr(view, 'datamodel', None): |
| continue |
| view.datamodel = CustomSQLAInterface(view.datamodel.obj) |
| |
| def init_role(self, role_name, role_vms, role_perms): |
| """ |
| Initialize the role with the permissions and related view-menus. |
| |
| :param role_name: |
| :param role_vms: |
| :param role_perms: |
| :return: |
| """ |
| pvms = self.get_session.query(sqla_models.PermissionView).all() |
| pvms = [p for p in pvms if p.permission and p.view_menu] |
| |
| role = self.find_role(role_name) |
| if not role: |
| role = self.add_role(role_name) |
| |
| if len(role.permissions) == 0: |
| self.log.info('Initializing permissions for role:%s in the database.', role_name) |
| role_pvms = set() |
| for pvm in pvms: |
| if pvm.view_menu.name in role_vms and pvm.permission.name in role_perms: |
| role_pvms.add(pvm) |
| role.permissions = list(role_pvms) |
| self.get_session.merge(role) |
| self.get_session.commit() |
| else: |
| self.log.debug('Existing permissions for the role:%s ' |
| 'within the database will persist.', role_name) |
| |
| def get_user_roles(self, user=None): |
| """ |
| Get all the roles associated with the user. |
| |
| :param user: the ab_user in FAB model. |
| :return: a list of roles associated with the user. |
| """ |
| if user is None: |
| user = g.user |
| if user.is_anonymous: |
| public_role = appbuilder.config.get('AUTH_ROLE_PUBLIC') |
| return [appbuilder.security_manager.find_role(public_role)] \ |
| if public_role else [] |
| return user.roles |
| |
| def get_all_permissions_views(self): |
| """ |
| Returns a set of tuples with the perm name and view menu name |
| """ |
| perms_views = set() |
| for role in self.get_user_roles(): |
| perms_views.update({(perm_view.permission.name, perm_view.view_menu.name) |
| for perm_view in role.permissions}) |
| return perms_views |
| |
| def get_accessible_dag_ids(self, username=None): |
| """ |
| Return a set of dags that user has access to(either read or write). |
| |
| :param username: Name of the user. |
| :return: A set of dag ids that the user could access. |
| """ |
| if not username: |
| username = g.user |
| |
| if username.is_anonymous or 'Public' in username.roles: |
| # return an empty set if the role is public |
| return set() |
| |
| roles = {role.name for role in username.roles} |
| if {'Admin', 'Viewer', 'User', 'Op'} & roles: |
| return self.DAG_VMS |
| |
| user_perms_views = self.get_all_permissions_views() |
| # return a set of all dags that the user could access |
| return set([view for perm, view in user_perms_views if perm in self.DAG_PERMS]) |
| |
| def has_access(self, permission, view_name, user=None): |
| """ |
| Verify whether a given user could perform certain permission |
| (e.g can_read, can_write) on the given dag_id. |
| |
| :param permission: permission on dag_id(e.g can_read, can_edit). |
| :type permission: str |
| :param view_name: name of view-menu(e.g dag id is a view-menu as well). |
| :type permission: str |
| :param user: user name |
| :type permission: str |
| :return: a bool whether user could perform certain permission on the dag_id. |
| :rtype bool |
| """ |
| if not user: |
| user = g.user |
| if user.is_anonymous: |
| return self.is_item_public(permission, view_name) |
| return self._has_view_access(user, permission, view_name) |
| |
| def _get_and_cache_perms(self): |
| """ |
| Cache permissions-views |
| """ |
| self.perms = self.get_all_permissions_views() |
| |
| def _has_role(self, role_name_or_list): |
| """ |
| Whether the user has this role name |
| """ |
| if not isinstance(role_name_or_list, list): |
| role_name_or_list = [role_name_or_list] |
| return any( |
| [r.name in role_name_or_list for r in self.get_user_roles()]) |
| |
| def _has_perm(self, permission_name, view_menu_name): |
| """ |
| Whether the user has this perm |
| """ |
| if hasattr(self, 'perms'): |
| if (permission_name, view_menu_name) in self.perms: |
| return True |
| # rebuild the permissions set |
| self._get_and_cache_perms() |
| return (permission_name, view_menu_name) in self.perms |
| |
| def has_all_dags_access(self): |
| """ |
| Has all the dag access in any of the 3 cases: |
| 1. Role needs to be in (Admin, Viewer, User, Op). |
| 2. Has can_dag_read permission on all_dags view. |
| 3. Has can_dag_edit permission on all_dags view. |
| """ |
| return ( |
| self._has_role(['Admin', 'Viewer', 'Op', 'User']) or |
| self._has_perm('can_dag_read', 'all_dags') or |
| self._has_perm('can_dag_edit', 'all_dags')) |
| |
| def clean_perms(self): |
| """ |
| FAB leaves faulty permissions that need to be cleaned up |
| """ |
| self.log.info('Cleaning faulty perms') |
| sesh = self.get_session |
| pvms = ( |
| sesh.query(sqla_models.PermissionView) |
| .filter(or_( |
| sqla_models.PermissionView.permission == None, # NOQA |
| sqla_models.PermissionView.view_menu == None, # NOQA |
| )) |
| ) |
| # Since FAB doesn't define ON DELETE CASCADE on these tables, we need |
| # to delete the _object_ so that SQLA knows to delete the many-to-many |
| # relationship object too. :( |
| |
| deleted_count = 0 |
| for pvm in pvms: |
| sesh.delete(pvm) |
| deleted_count += 1 |
| sesh.commit() |
| if deleted_count: |
| self.log.info('Deleted %s faulty permissions', deleted_count) |
| |
| def _merge_perm(self, permission_name, view_menu_name): |
| """ |
| Add the new permission , view_menu to ab_permission_view_role if not exists. |
| It will add the related entry to ab_permission |
| and ab_view_menu two meta tables as well. |
| |
| :param permission_name: Name of the permission. |
| :type permission_name: str |
| :param view_menu_name: Name of the view-menu |
| :type view_menu_name: str |
| :return: |
| """ |
| permission = self.find_permission(permission_name) |
| view_menu = self.find_view_menu(view_menu_name) |
| pv = None |
| if permission and view_menu: |
| pv = self.get_session.query(self.permissionview_model).filter_by( |
| permission=permission, view_menu=view_menu).first() |
| if not pv and permission_name and view_menu_name: |
| self.add_permission_view_menu(permission_name, view_menu_name) |
| |
| @provide_session |
| def create_custom_dag_permission_view(self, session=None): |
| """ |
| Workflow: |
| 1. Fetch all the existing (permissions, view-menu) from Airflow DB. |
| 2. Fetch all the existing dag models that are either active or paused. |
| 3. Create both read and write permission view-menus relation for every dags from step 2 |
| 4. Find out all the dag specific roles(excluded pubic, admin, viewer, op, user) |
| 5. Get all the permission-vm owned by the user role. |
| 6. Grant all the user role's permission-vm except the all-dag view-menus to the dag roles. |
| 7. Commit the updated permission-vm-role into db |
| |
| :return: None. |
| """ |
| # todo(Tao): should we put this function here or in scheduler loop? |
| self.log.info('Fetching a set of all permission, view_menu from FAB meta-table') |
| |
| def merge_pv(perm, view_menu): |
| """Create permission view menu only if it doesn't exist""" |
| if view_menu and perm and (view_menu, perm) not in all_pvs: |
| self._merge_perm(perm, view_menu) |
| |
| all_pvs = set() |
| for pv in self.get_session.query(self.permissionview_model).all(): |
| if pv.permission and pv.view_menu: |
| all_pvs.add((pv.permission.name, pv.view_menu.name)) |
| |
| # Get all the active / paused dags and insert them into a set |
| all_dags_models = session.query(models.DagModel)\ |
| .filter(or_(models.DagModel.is_active, models.DagModel.is_paused)).all() |
| |
| # create can_dag_edit and can_dag_read permissions for every dag(vm) |
| for dag in all_dags_models: |
| for perm in self.DAG_PERMS: |
| merge_pv(perm, dag.dag_id) |
| |
| # for all the dag-level role, add the permission of viewer |
| # with the dag view to ab_permission_view |
| all_roles = self.get_all_roles() |
| user_role = self.find_role('User') |
| |
| dag_role = [role for role in all_roles if role.name not in EXISTING_ROLES] |
| update_perm_views = [] |
| |
| # need to remove all_dag vm from all the existing view-menus |
| dag_vm = self.find_view_menu('all_dags') |
| ab_perm_view_role = sqla_models.assoc_permissionview_role |
| perm_view = self.permissionview_model |
| view_menu = self.viewmenu_model |
| |
| all_perm_view_by_user = session.query(ab_perm_view_role)\ |
| .join(perm_view, perm_view.id == ab_perm_view_role |
| .columns.permission_view_id)\ |
| .filter(ab_perm_view_role.columns.role_id == user_role.id)\ |
| .join(view_menu)\ |
| .filter(perm_view.view_menu_id != dag_vm.id) |
| all_perm_views = set([role.permission_view_id for role in all_perm_view_by_user]) |
| |
| for role in dag_role: |
| # Get all the perm-view of the role |
| existing_perm_view_by_user = self.get_session.query(ab_perm_view_role)\ |
| .filter(ab_perm_view_role.columns.role_id == role.id) |
| |
| existing_perms_views = set([pv.permission_view_id |
| for pv in existing_perm_view_by_user]) |
| missing_perm_views = all_perm_views - existing_perms_views |
| |
| for perm_view_id in missing_perm_views: |
| update_perm_views.append({'permission_view_id': perm_view_id, |
| 'role_id': role.id}) |
| |
| if update_perm_views: |
| self.get_session.execute(ab_perm_view_role.insert(), update_perm_views) |
| self.get_session.commit() |
| |
| def update_admin_perm_view(self): |
| """ |
| Admin should has all the permission-views, except the dag views. |
| because Admin have already have all_dags permission. |
| Add the missing ones to the table for admin. |
| |
| :return: None. |
| """ |
| all_dag_view = self.find_view_menu('all_dags') |
| dag_perm_ids = [self.find_permission('can_dag_edit').id, self.find_permission('can_dag_read').id] |
| pvms = self.get_session.query(sqla_models.PermissionView).filter(~and_( |
| sqla_models.PermissionView.permission_id.in_(dag_perm_ids), |
| sqla_models.PermissionView.view_menu_id != all_dag_view.id) |
| ).all() |
| |
| pvms = [p for p in pvms if p.permission and p.view_menu] |
| |
| admin = self.find_role('Admin') |
| admin.permissions = list(set(admin.permissions) | set(pvms)) |
| |
| self.get_session.commit() |
| |
| def sync_roles(self): |
| """ |
| 1. Init the default role(Admin, Viewer, User, Op, public) |
| with related permissions. |
| 2. Init the custom role(dag-user) with related permissions. |
| |
| :return: None. |
| """ |
| self.log.info('Start syncing user roles.') |
| # Create global all-dag VM |
| self.create_perm_vm_for_all_dag() |
| |
| # Create default user role. |
| for config in self.ROLE_CONFIGS: |
| role = config['role'] |
| vms = config['vms'] |
| perms = config['perms'] |
| self.init_role(role, vms, perms) |
| self.create_custom_dag_permission_view() |
| |
| # init existing roles, the rest role could be created through UI. |
| self.update_admin_perm_view() |
| self.clean_perms() |
| |
| def sync_perm_for_dag(self, dag_id, access_control=None): |
| """ |
| Sync permissions for given dag id. The dag id surely exists in our dag bag |
| as only / refresh button or cli.sync_perm will call this function |
| |
| :param dag_id: the ID of the DAG whose permissions should be updated |
| :type dag_id: str |
| :param access_control: a dict where each key is a rolename and |
| each value is a set() of permission names (e.g., |
| {'can_dag_read'} |
| :type access_control: dict |
| :return: |
| """ |
| for dag_perm in self.DAG_PERMS: |
| perm_on_dag = self.find_permission_view_menu(dag_perm, dag_id) |
| if perm_on_dag is None: |
| self.add_permission_view_menu(dag_perm, dag_id) |
| |
| if access_control: |
| self._sync_dag_view_permissions(dag_id, access_control) |
| |
| def _sync_dag_view_permissions(self, dag_id, access_control): |
| """Set the access policy on the given DAG's ViewModel. |
| |
| :param dag_id: the ID of the DAG whose permissions should be updated |
| :type dag_id: str |
| :param access_control: a dict where each key is a rolename and |
| each value is a set() of permission names (e.g., |
| {'can_dag_read'} |
| :type access_control: dict |
| """ |
| def _get_or_create_dag_permission(perm_name): |
| dag_perm = self.find_permission_view_menu(perm_name, dag_id) |
| if not dag_perm: |
| self.log.info("Creating new permission '{}' on view '{}'".format( |
| perm_name, |
| dag_id |
| )) |
| dag_perm = self.add_permission_view_menu(perm_name, dag_id) |
| |
| return dag_perm |
| |
| def _revoke_stale_permissions(dag_view): |
| existing_dag_perms = self.find_permissions_view_menu(dag_view) |
| for perm in existing_dag_perms: |
| non_admin_roles = [role for role in perm.role |
| if role.name != 'Admin'] |
| for role in non_admin_roles: |
| target_perms_for_role = access_control.get(role.name, {}) |
| if perm.permission.name not in target_perms_for_role: |
| self.log.info("Revoking '{}' on DAG '{}' for role '{}'".format( |
| perm.permission, |
| dag_id, |
| role.name |
| )) |
| self.del_permission_role(role, perm) |
| |
| dag_view = self.find_view_menu(dag_id) |
| if dag_view: |
| _revoke_stale_permissions(dag_view) |
| |
| for rolename, perms in access_control.items(): |
| role = self.find_role(rolename) |
| if not role: |
| raise AirflowException( |
| "The access_control mapping for DAG '{}' includes a role " |
| "named '{}', but that role does not exist".format( |
| dag_id, |
| rolename)) |
| |
| perms = set(perms) |
| invalid_perms = perms - self.DAG_PERMS |
| if invalid_perms: |
| raise AirflowException( |
| "The access_control map for DAG '{}' includes the following " |
| "invalid permissions: {}; The set of valid permissions " |
| "is: {}".format(dag_id, |
| (perms - self.DAG_PERMS), |
| self.DAG_PERMS)) |
| |
| for perm_name in perms: |
| dag_perm = _get_or_create_dag_permission(perm_name) |
| self.add_permission_role(role, dag_perm) |
| |
| def create_perm_vm_for_all_dag(self): |
| """ |
| Create perm-vm if not exist and insert into FAB security model for all-dags. |
| """ |
| # create perm for global logical dag |
| for dag_vm in self.DAG_VMS: |
| for perm in self.DAG_PERMS: |
| self._merge_perm(permission_name=perm, |
| view_menu_name=dag_vm) |
