fix: Addressing the issue of unauthorized execution of system tools (#14209)

* fix: Addressing the issue of unauthorized execution of system tools

* perf: Optimization conditions

---------

Co-authored-by: jiangweidong <1053570670@qq.com>
pull/14249/head
fit2bot 2 months ago committed by GitHub
parent 0311446384
commit 43fa3f420a
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

@ -3,10 +3,11 @@
import json import json
import asyncio import asyncio
from asgiref.sync import sync_to_async
from channels.generic.websocket import AsyncJsonWebsocketConsumer from channels.generic.websocket import AsyncJsonWebsocketConsumer
from django.core.cache import cache from django.core.cache import cache
from django.conf import settings from django.conf import settings
from django.utils.translation import gettext_lazy as _, activate from django.utils.translation import gettext_lazy as _
from django.utils import translation from django.utils import translation
from urllib.parse import parse_qs from urllib.parse import parse_qs
@ -37,13 +38,29 @@ TASK_STATUS_IS_OVER = 'OVER'
class ToolsWebsocket(AsyncJsonWebsocketConsumer): class ToolsWebsocket(AsyncJsonWebsocketConsumer):
is_closed: bool = False
@staticmethod
@sync_to_async
def get_user_roles(user):
return [str(i) for i in user.system_roles.values_list('id', flat=True)]
async def is_superuser(self, user):
from rbac.builtin import BuiltinRole
ids = await self.get_user_roles(user)
return BuiltinRole.system_admin.id in ids
async def connect(self): async def connect(self):
user = self.scope["user"] user = self.scope["user"]
if user.is_authenticated: if user.is_authenticated:
has_perm = await sync_to_async(user.has_perm)('rbac.view_systemtools')
if await self.is_superuser(user) or (settings.TOOL_USER_ENABLED and has_perm):
await self.accept() await self.accept()
else: else:
await self.close() await self.close()
else:
await self.close()
async def send_msg(self, msg=''): async def send_msg(self, msg=''):
await self.send_json({'msg': f'{msg}\r\n'}) await self.send_json({'msg': f'{msg}\r\n'})
@ -96,6 +113,12 @@ class ToolsWebsocket(AsyncJsonWebsocketConsumer):
await self.send_msg() await self.send_msg()
await self.close() await self.close()
async def close(self, code=None):
if self.is_closed:
return
await super().close(code)
self.is_closed = True
async def disconnect(self, code): async def disconnect(self, code):
await self.close() await self.close()
close_old_connections() close_old_connections()

Loading…
Cancel
Save