def ready(self): from orgs.models import Organization from orgs.utils import set_current_org set_current_org(Organization.root()) super().ready() from .celery import signal_handler
def setUp(self): Organization.objects.bulk_create([ Organization(name='org-01'), Organization(name='org-02'), Organization(name='org-03'), ]) org_01, org_02, org_03 = Organization.objects.all() self.org_01, self.org_02, self.org_03 = org_01, org_02, org_03 set_current_org(org_01) AdminUser.objects.bulk_create([ AdminUser(name='au-01', username='******'), AdminUser(name='au-02', username='******'), AdminUser(name='au-03', username='******'), ]) SystemUser.objects.bulk_create([ SystemUser(name='su-01', username='******'), SystemUser(name='su-02', username='******'), SystemUser(name='su-03', username='******'), ]) admin_users = AdminUser.objects.all() Asset.objects.bulk_create([ Asset(hostname='asset-01', ip='192.168.1.1', public_ip='192.168.1.1', admin_user=admin_users[0]), Asset(hostname='asset-02', ip='192.168.1.2', public_ip='192.168.1.2', admin_user=admin_users[0]), Asset(hostname='asset-03', ip='192.168.1.3', public_ip='192.168.1.3', admin_user=admin_users[0]), ]) new_user = User.objects.create new_org_memeber = OrganizationMember.objects.create u = new_user(name='user-01', username='******', email='*****@*****.**') new_org_memeber(org=org_01, user=u, role=ORG_ROLE.USER) new_org_memeber(org=org_02, user=u, role=ORG_ROLE.USER) self.user_01 = u u = new_user(name='org-admin-01', username='******', email='*****@*****.**') new_org_memeber(org=org_01, user=u, role=ORG_ROLE.ADMIN) self.org_admin_01 = u u = new_user(name='org-admin-02', username='******', email='*****@*****.**') new_org_memeber(org=org_02, user=u, role=ORG_ROLE.ADMIN) self.org_admin_02 = u
def finalize_response(self, request, response, *args, **kwargs): response = super().finalize_response(request, response, *args, **kwargs) org = getattr(self, 'current_org', None) if org: set_current_org(org) return response
def create_root_node(cls): # 如果使用current_org 在set_current_org时会死循环 _current_org = get_current_org() with transaction.atomic(): if not _current_org.is_real(): return cls.default_node() set_current_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) or ['1'] key = max([int(k) for k in org_nodes_roots_keys]) key = str(key + 1) if key != 0 else '2' set_current_org(_current_org) root = cls.objects.create(key=key, value=_current_org.name) return root
def create_root_node(cls): # 如果使用current_org 在set_current_org时会死循环 _current_org = get_current_org() with transaction.atomic(): if _current_org.is_default(): key = '0' else: set_current_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) or [0] key = max([int(k) for k in org_nodes_roots_keys]) + 1 set_current_org(_current_org) root = cls.objects.create(key=key, value=_current_org.name) return root
def change_key0_to_key1(apps, schema_editor): from orgs.utils import set_current_org # https://stackoverflow.com/questions/28777338/django-migrations-runpython-not-able-to-call-model-methods Organization = apps.get_model('orgs', 'Organization') Node = apps.get_model('assets', 'Node') print() org = Organization.objects.get(id=default_id) set_current_org(org) exists_0 = Node.objects.filter(key__startswith='0').exists() if not exists_0: print(f'--> Not exist key=0 nodes, do nothing.') return key_1_count = Node.objects.filter(key__startswith='1').count() if key_1_count > 1: print( f'--> Node key=1 have children, can`t just delete it. Please contact JumpServer team' ) return root_node = Node.objects.filter(key='1').first() if root_node and root_node.assets.exists(): print(f'--> Node key=1 has assets, do nothing.') return with atomic(): if root_node: print(f'--> Delete node key=1') root_node.delete() nodes_0 = Node.objects.filter(key__startswith='0') for n in nodes_0: old_key = n.key key_list = n.key.split(':') key_list[0] = '1' new_key = ':'.join(key_list) new_parent_key = ':'.join(key_list[:-1]) n.key = new_key n.parent_key = new_parent_key n.save() print('--> Modify key ( {} > {} )'.format(old_key, new_key))
def start(self): self.task.latest_history = self self.task.save() current_org = get_current_org() set_to_root_org() time_start = time.time() date_start = timezone.now() is_success = False summary = {} raw = '' try: date_start_s = datetime.datetime.now().strftime( '%Y-%m-%d %H:%M:%S') print(_("{} Start task: {}").format(date_start_s, self.task.name)) raw, summary = self.start_runner() is_success = summary.get('success', False) except Exception as e: logger.error(e, exc_info=True) raw = {"dark": {"all": str(e)}, "contacted": []} finally: date_end = timezone.now() date_end_s = date_end.strftime('%Y-%m-%d %H:%M:%S') print(_("{} Task finish").format(date_end_s)) print('.\n\n.') task = Task.objects.get(id=self.task_id) task.total_run_amount = models.F('total_run_amount') + 1 if is_success: task.success_run_amount = models.F('success_run_amount') + 1 task.save() AdHocRunHistory.objects.filter(id=self.id).update( date_start=date_start, is_finished=True, is_success=is_success, date_finished=timezone.now(), timedelta=time.time() - time_start, summary=summary) set_current_org(current_org) return raw, summary