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 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 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
