From b6ab3df038af96b7ea0a7c8443d656a8c9825051 Mon Sep 17 00:00:00 2001 From: ibuler Date: Mon, 18 Mar 2024 18:53:22 +0800 Subject: [PATCH] =?UTF-8?q?perf:=20=E4=BC=98=E5=8C=96=20celery=20task=20lo?= =?UTF-8?q?g=20=E6=9D=83=E9=99=90=E6=8E=A7=E5=88=B6?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../rewriting/storage/permissions.py | 2 +- .../0028_celerytaskexecution_creator.py | 21 +++++++++++++++ apps/ops/models/celery.py | 2 ++ apps/ops/signal_handlers.py | 8 ++++-- apps/ops/ws.py | 27 +++++++++++++++++-- 5 files changed, 55 insertions(+), 5 deletions(-) create mode 100644 apps/ops/migrations/0028_celerytaskexecution_creator.py diff --git a/apps/jumpserver/rewriting/storage/permissions.py b/apps/jumpserver/rewriting/storage/permissions.py index 7d91c246f..c6473df19 100644 --- a/apps/jumpserver/rewriting/storage/permissions.py +++ b/apps/jumpserver/rewriting/storage/permissions.py @@ -4,7 +4,7 @@ path_perms_map = { 'xpack': '*', 'settings': '*', 'img': '*', - 'replay': 'default', + 'replay': 'terminal.view_sessionreplay', 'applets': 'terminal.view_applet', 'virtual_apps': 'terminal.view_virtualapp', 'playbooks': 'ops.view_playbook' diff --git a/apps/ops/migrations/0028_celerytaskexecution_creator.py b/apps/ops/migrations/0028_celerytaskexecution_creator.py new file mode 100644 index 000000000..37d44077b --- /dev/null +++ b/apps/ops/migrations/0028_celerytaskexecution_creator.py @@ -0,0 +1,21 @@ +# Generated by Django 4.1.13 on 2024-03-18 06:47 + +from django.conf import settings +from django.db import migrations, models +import django.db.models.deletion + + +class Migration(migrations.Migration): + + dependencies = [ + migrations.swappable_dependency(settings.AUTH_USER_MODEL), + ('ops', '0027_alter_celerytaskexecution_options'), + ] + + operations = [ + migrations.AddField( + model_name='celerytaskexecution', + name='creator', + field=models.ForeignKey(db_constraint=False, default=None, null=True, on_delete=django.db.models.deletion.SET_NULL, to=settings.AUTH_USER_MODEL, verbose_name='Creator'), + ), + ] diff --git a/apps/ops/models/celery.py b/apps/ops/models/celery.py index ef1fab463..9e40cd921 100644 --- a/apps/ops/models/celery.py +++ b/apps/ops/models/celery.py @@ -82,6 +82,8 @@ class CeleryTaskExecution(models.Model): kwargs = models.JSONField(verbose_name=_("Kwargs")) state = models.CharField(max_length=16, verbose_name=_("State")) is_finished = models.BooleanField(default=False, verbose_name=_("Finished")) + creator = models.ForeignKey('users.User', on_delete=models.SET_NULL, default=None, null=True, + verbose_name=_('Creator'), db_constraint=False) date_published = models.DateTimeField(auto_now_add=True, verbose_name=_('Date published')) date_start = models.DateTimeField(null=True, verbose_name=_('Date start')) date_finished = models.DateTimeField(null=True, verbose_name=_('Date finished')) diff --git a/apps/ops/signal_handlers.py b/apps/ops/signal_handlers.py index d00d33ad0..5df6f6cf3 100644 --- a/apps/ops/signal_handlers.py +++ b/apps/ops/signal_handlers.py @@ -1,8 +1,7 @@ import ast -import psutil -from psutil import NoSuchProcess import time +import psutil from celery import signals from django.core.cache import cache from django.db import transaction @@ -11,10 +10,12 @@ from django.db.utils import ProgrammingError from django.dispatch import receiver from django.utils import translation, timezone from django.utils.functional import LazyObject +from psutil import NoSuchProcess from common.db.utils import close_old_connections, get_logger from common.signals import django_ready from common.utils.connection import RedisPubSub +from jumpserver.utils import get_current_request from orgs.utils import get_current_org_id, set_current_org from .celery import app from .models import CeleryTaskExecution, CeleryTask, Job @@ -146,6 +147,9 @@ def task_sent_handler(headers=None, body=None, **kwargs): 'args': args, 'kwargs': kwargs } + request = get_current_request() + if request and request.user.is_authenticated: + data['creator'] = request.user CeleryTaskExecution.objects.create(**data) CeleryTask.objects.filter(name=task).update(date_last_publish=timezone.now()) diff --git a/apps/ops/ws.py b/apps/ops/ws.py index c2386c77a..285503f6a 100644 --- a/apps/ops/ws.py +++ b/apps/ops/ws.py @@ -2,6 +2,7 @@ import asyncio import os import aiofiles +from asgiref.sync import sync_to_async from channels.generic.websocket import AsyncJsonWebsocketConsumer from common.db.utils import close_old_connections @@ -9,12 +10,17 @@ from common.utils import get_logger from .ansible.utils import get_ansible_task_log_path from .celery.utils import get_celery_task_log_path from .const import CELERY_LOG_MAGIC_MARK +from .models import CeleryTaskExecution logger = get_logger(__name__) class TaskLogWebsocket(AsyncJsonWebsocketConsumer): disconnected = False + user_tasks = ( + 'ops.tasks.run_ops_job', + 'ops.tasks.run_ops_job_execution', + ) log_types = { 'celery': get_celery_task_log_path, @@ -33,10 +39,27 @@ class TaskLogWebsocket(AsyncJsonWebsocketConsumer): if func: return func(task_id) + @sync_to_async + def get_task(self, task_id): + task = CeleryTaskExecution.objects.filter(id=task_id).first() + # task.creator 是 foreign key, 会异步去查询的,在下面的 if task.creator 会报错, 所以这里先取出来 + if task and task.creator != ' ': + return task + else: + return None + async def receive_json(self, content, **kwargs): task_id = content.get('task') - task_typ = content.get('type', 'celery') - log_path = self.get_log_path(task_id, task_typ) + task = await self.get_task(task_id) + if not task: + await self.send_json({'message': 'Task not found', 'task': task_id}) + return + if task.name in self.user_tasks and task.creator != self.scope['user']: + await self.send_json({'message': 'No permission', 'task': task_id}) + return + + task_type = content.get('type', 'celery') + log_path = self.get_log_path(task_id, task_type) await self.async_handle_task(task_id, log_path) async def async_handle_task(self, task_id, log_path):