def modify_other_org_root_node_key(cls): """ 解决创建 default 节点失败的问题, 因为在其他组织下存在 default 节点,故在 DEFAULT 组织下 get 不到 create 失败 """ logger.info("Modify other org root node key") with tmp_to_org(Organization.root()): node_key1 = cls.objects.filter(key='1').first() if not node_key1: logger.info("Not found node that `key` = 1") return if not node_key1.org.is_real(): logger.info("Org is not real for node that `key` = 1") return with transaction.atomic(): with tmp_to_org(node_key1.org): org_root_node_new_key = cls.get_next_org_root_node_key() for n in cls.objects.all(): old_key = n.key key_list = n.key.split(':') key_list[0] = org_root_node_new_key new_key = ':'.join(key_list) n.key = new_key n.save() logger.info('Modify key ( {} > {} )'.format(old_key, new_key))
def get_org_related_flows(cls, org_id=None): if org_id: with tmp_to_org(org_id): flows = cls.objects.all() else: flows = cls.objects.all() cur_flow_types = flows.values_list('type', flat=True) root_id = Organization.ROOT_ID with tmp_to_org(root_id): diff_global_flows = cls.objects.exclude(type__in=cur_flow_types).filter(org_id=root_id) return flows | diff_global_flows
def refresh_if_need(self, force=False): user = self.user with tmp_to_root_org(): UserAssetGrantedTreeNodeRelation.objects.filter(user=user).exclude( org_id__in=self.org_ids).delete() if force or self.have_need_refresh_orgs(): with UserGrantedTreeRebuildLock(user_id=user.id): if force: orgs = self.orgs self.set_all_orgs_as_builed() else: orgs = self.get_need_refresh_orgs_and_fill_up() for org in orgs: with tmp_to_org(org): t_start = time.time() logger.info( f'Rebuild user tree: user={self.user} org={current_org}' ) utils = UserGrantedTreeBuildUtils(user) utils.rebuild_user_granted_tree() logger.info( f'Rebuild user tree ok: cost={time.time() - t_start} user={self.user} org={current_org}' )
def default_node(cls): default_org = Organization.default() with tmp_to_org(default_org): defaults = {'value': default_org.name} obj, created = cls.objects.get_or_create(defaults=defaults, key=cls.default_key) return obj
def on_user_created_set_default_org(sender, instance, created, **kwargs): if not created: return if instance.orgs.count() > 0: return with tmp_to_org(Organization.default()): Organization.default().add_member(instance)
def post(self, request, *args, **kwargs): serializer = self.serializer_class(data=request.data) if not serializer.is_valid(): msg = str(serializer.errors) return Response({'ok': False, 'msg': msg}, status=401) user_id = serializer.validated_data['user_id'] session_id = serializer.validated_data['session_id'] with tmp_to_root_org(): session = get_object_or_none(Session, pk=session_id) if not session: msg = _('Session does not exist: {}'.format(session_id)) return Response({'ok': False, 'msg': msg}, status=401) if not session.can_join: msg = _('Session is finished or the protocol not supported') return Response({'ok': False, 'msg': msg}, status=401) user = get_object_or_none(User, pk=user_id) if not user: msg = _('User does not exist: {}'.format(user_id)) return Response({'ok': False, 'msg': msg}, status=401) with tmp_to_org(session.org): if not user.admin_or_audit_orgs: msg = _('User does not have permission') return Response({'ok': False, 'msg': msg}, status=401) return Response({'ok': True, 'msg': ''}, status=200)
def on_node_asset_change(sender, action, instance, reverse, pk_set, **kwargs): # 不允许 `pre_clear` ,因为该信号没有 `pk_set` # [官网](https://docs.djangoproject.com/en/3.1/ref/signals/#m2m-changed) refused = (PRE_CLEAR, ) if action in refused: raise ValueError mapper = {PRE_ADD: add, POST_REMOVE: sub} if action not in mapper: return operator = mapper[action] with tmp_to_org(instance.org): if reverse: node: Node = instance asset_pk_set = set(pk_set) NodeAssetsAmountUtils.update_node_assets_amount( node, asset_pk_set, operator) else: asset_pk = instance.id # 与资产直接关联的节点 node_keys = set( Node.objects.filter(id__in=pk_set).values_list('key', flat=True)) NodeAssetsAmountUtils.update_nodes_asset_amount( node_keys, asset_pk, operator)
def _construct_meta_display_of_approve(self): meta_display_fields = [ 'approve_actions_display', 'approve_assets_display', 'approve_system_users_display' ] approve_actions = self.ticket.meta.get('approve_actions', Action.NONE) approve_actions_display = Action.value_to_choices_display( approve_actions) approve_asset_ids = self.ticket.meta.get('approve_assets', []) approve_system_user_ids = self.ticket.meta.get('approve_system_users', []) with tmp_to_org(self.ticket.org_id): assets = Asset.objects.filter(id__in=approve_asset_ids) system_users = SystemUser.objects.filter( id__in=approve_system_user_ids) approve_assets_display = [str(asset) for asset in assets] approve_system_users_display = [ str(system_user) for system_user in system_users ] meta_display_values = [ approve_actions_display, approve_assets_display, approve_system_users_display ] meta_display = dict(zip(meta_display_fields, meta_display_values)) return meta_display
def favorite_node(cls): with tmp_to_org(Organization.system()): defaults = {'value': cls.favorite_value} obj, created = cls.objects.get_or_create( defaults=defaults, key=cls.favorite_key ) return obj
def add_asset_users(assets, results): assets_map = {a.hostname: a for a in assets} parser_map = { 'linux': parse_linux_result_to_users, 'windows': parse_windows_result_to_users } assets_users_map = {} for platform, platform_results in results.items(): for hostname, res in platform_results.items(): parse = parser_map.get(platform) users = parse(res) logger.debug('Gathered host users: {} {}'.format(hostname, users)) asset = assets_map.get(hostname) if not asset: continue assets_users_map[asset] = users for asset, users in assets_users_map.items(): with tmp_to_org(asset.org_id): GatheredUser.objects.filter(asset=asset, present=True)\ .update(present=False) for username, data in users.items(): defaults = {'asset': asset, 'username': username, 'present': True} if data.get("ip"): defaults["ip_last_login"] = data["ip"] if data.get("date"): defaults["date_last_login"] = data["date"] GatheredUser.objects.update_or_create( defaults=defaults, asset=asset, username=username, )
def on_org_delete(sender, instance, **kwargs): expire_orgs_mapping_for_memory(instance.id) # 删除该组织下所有 节点 with tmp_to_org(instance): root_node = Node.org_root() if root_node: root_node.delete()
def add_need_refresh_by_asset_perm_ids_cross_orgs(cls, asset_perm_ids): org_id_perm_ids_mapper = defaultdict(set) pairs = AssetPermission.objects.filter(id__in=asset_perm_ids).values_list('org_id', 'id') for org_id, perm_id in pairs: org_id_perm_ids_mapper[org_id].add(perm_id) for org_id, perm_ids in org_id_perm_ids_mapper.items(): with tmp_to_org(org_id): cls.add_need_refresh_by_asset_perm_ids(perm_ids)
def create(self, validated_data): instance = super().create(validated_data) name = _('Created by ticket ({}-{})').format(instance.title, str(instance.id)[:4]) with tmp_to_org(instance.org_id): if not AssetPermission.objects.filter(name=name).exists(): instance.meta.update({'apply_permission_name': name}) return instance raise serializers.ValidationError(_('Permission named `{}` already exists'.format(name)))
def execute_account_backup_plan(pid, trigger): with tmp_to_root_org(): plan = get_object_or_none(AccountBackupPlan, pk=pid) if not plan: logger.error("No account backup route plan found: {}".format(pid)) return with tmp_to_org(plan.org): plan.execute(trigger)
def validate_apply_applications(self, apply_applications): type = self.root.initial_data['meta'].get('apply_type') org_id = self.root.initial_data.get('org_id') with tmp_to_org(org_id): applications = Application.objects.filter( id__in=apply_applications, type=type).values_list('id', flat=True) return list(applications)
def on_permission_assets_changed(sender, instance, action, reverse, pk_set, model, **kwargs): if reverse: raise M2MReverseNotAllowed if not need_rebuild_mapping_node(action): return with tmp_to_org(instance.org): UserGrantedTreeRefreshController.add_need_refresh_by_asset_perm_ids([instance.id])
def check_node_assets_amount_task(org_id=Organization.ROOT_ID): try: with tmp_to_org(Organization.get_instance(org_id)): check_node_assets_amount() except AcquireFailed: logger.error( _('The task of self-checking is already running and cannot be started repeatedly' ))
def default_node(cls): with tmp_to_org(Organization.default()): defaults = {'value': cls.default_value} obj, created = cls.objects.get_or_create( defaults=defaults, key=cls.default_key, ) return obj
def get_next_org_root_node_key(cls): with tmp_to_org(Organization.root()): org_nodes_roots = cls.objects.filter(key__regex=r'^[0-9]+$') org_nodes_roots_keys = org_nodes_roots.values_list('key', flat=True) if not org_nodes_roots_keys: org_nodes_roots_keys = ['1'] max_key = max([int(k) for k in org_nodes_roots_keys]) key = str(max_key + 1) if max_key != 0 else '2' return key
def on_asset_perm_pre_save(sender, instance, **kwargs): try: old = AssetPermission.objects.get(id=instance.id) if old.is_valid != instance.is_valid: with tmp_to_org(instance.org): UserGrantedTreeRefreshController.add_need_refresh_by_asset_perm_ids([instance.id]) except AssetPermission.DoesNotExist: pass
def new(cls, org_id=None): from ..utils import TreeService logger.debug("Create node tree") if not org_id: org_id = current_org.id with tmp_to_org(org_id): tree = TreeService.new() obj = cls(tree, org_id) obj.tree = tree return obj
def get_org_related_flows(cls, org_id=None): if org_id: with tmp_to_org(org_id): flows = cls.objects.all() else: flows = cls.objects.all() cur_flow_types = flows.values_list('type', flat=True) with tmp_to_root_org(): diff_global_flows = cls.objects.exclude(type__in=cur_flow_types) return flows | diff_global_flows
def get_object(self): instance = super().get_object() username = instance.username if instance.username_same_with_user: username = self.request.query_params.get("username") asset_id = self.kwargs.get('aid') asset = get_object_or_404(Asset, pk=asset_id) with tmp_to_org(asset.org_id): instance.load_asset_special_auth(asset=asset, username=username) return instance
def on_asset_permission_users_changed(sender, action, reverse, instance, pk_set, **kwargs): if reverse: raise M2MReverseNotAllowed if not need_rebuild_mapping_node(action): return with tmp_to_org(instance.org): UserGrantedTreeRefreshController.add_need_refresh_orgs_for_users( [current_org.id], pk_set )
def validate_approve_permission_name(self, permission_name): if not isinstance(self.root.instance, Ticket): return permission_name with tmp_to_org(self.root.instance.org_id): already_exists = AssetPermission.objects.filter( name=permission_name).exists() if not already_exists: return permission_name raise serializers.ValidationError( _('Permission named `{}` already exists'.format(permission_name)))
def _construct_meta_display_of_approve(self): meta_display_fields = ['approve_applications_display', 'approve_system_users_display'] approve_applications_id = self.ticket.meta.get('approve_applications', []) approve_system_users_id = self.ticket.meta.get('approve_system_users', []) with tmp_to_org(self.ticket.org_id): approve_applications = Application.objects.filter(id__in=approve_applications_id) system_users = SystemUser.objects.filter(id__in=approve_system_users_id) approve_applications_display = [str(application) for application in approve_applications] approve_system_users_display = [str(system_user) for system_user in system_users] meta_display_values = [approve_applications_display, approve_system_users_display] meta_display = dict(zip(meta_display_fields, meta_display_values)) return meta_display
def test_system_user_connectivity_period(): if not const.PERIOD_TASK_ENABLED: logger.debug( "Period task disabled, test system user connectivity pass") return queryset_map = SystemUser.objects.all_group_by_org() for org, system_user in queryset_map.items(): task_name = _("Test system user connectivity period: {}").format( system_user) with tmp_to_org(org): assets = system_user.get_related_assets() test_system_user_connectivity_util(system_user, assets, task_name)
def get_assignees(self, org_id=None): assignees = [] org_id = org_id if org_id else get_current_org_id() with tmp_to_org(org_id): if self.strategy == TicketApprovalStrategy.super_admin: assignees = User.get_super_admins() elif self.strategy == TicketApprovalStrategy.org_admin: assignees = User.get_org_admins() elif self.strategy == TicketApprovalStrategy.super_org_admin: assignees = User.get_super_and_org_admins() elif self.strategy == TicketApprovalStrategy.custom_user: assignees = self.assignees.all() return assignees
def on_node_asset_change(action, instance, reverse, pk_set, **kwargs): if not need_rebuild_mapping_node(action): return if reverse: asset_pk_set = pk_set node_pk_set = [instance.id] else: asset_pk_set = [instance.id] node_pk_set = pk_set with tmp_to_org(instance.org): UserGrantedTreeRefreshController.add_need_refresh_on_nodes_assets_relate_change(node_pk_set, asset_pk_set)
def default_node(cls): with tmp_to_org(Organization.default()): defaults = {'value': cls.default_value} try: obj, created = cls.objects.get_or_create( defaults=defaults, key=cls.default_key, ) except IntegrityError as e: logger.error("Create default node failed: {}".format(e)) cls.modify_other_org_root_node_key() obj, created = cls.objects.get_or_create( defaults=defaults, key=cls.default_key, ) return obj