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
