From c44206cde694c8af027878aeef8c3f5db5ea2983 Mon Sep 17 00:00:00 2001 From: Tao Chen Date: Thu, 29 Aug 2024 00:21:47 +0800 Subject: [PATCH] users' devices will add to peers automatically --- .gitignore | 5 +- api/migrations/0008_rustdesdevice_owner.py | 18 +++ api/models_user.py | 104 +++++++++++++---- api/models_work.py | 127 +++++++++++++++++---- api/views_api.py | 114 ++++++++++++++++-- api/views_front.py | 1 - db/db.sqlite3 | Bin 163840 -> 167936 bytes rustdesk_server_api/settings.py | 4 + 8 files changed, 318 insertions(+), 55 deletions(-) create mode 100644 api/migrations/0008_rustdesdevice_owner.py diff --git a/.gitignore b/.gitignore index fea7e31..e47103c 100644 --- a/.gitignore +++ b/.gitignore @@ -28,4 +28,7 @@ LICENSE.rst db/test_db.sqlite3 job2en.py -新建文本文档.txt \ No newline at end of file +新建文本文档.txt + +venv +.env \ No newline at end of file diff --git a/api/migrations/0008_rustdesdevice_owner.py b/api/migrations/0008_rustdesdevice_owner.py new file mode 100644 index 0000000..1736ce1 --- /dev/null +++ b/api/migrations/0008_rustdesdevice_owner.py @@ -0,0 +1,18 @@ +# Generated by Django 5.1 on 2024-08-29 00:01 + +from django.db import migrations, models + + +class Migration(migrations.Migration): + + dependencies = [ + ('api', '0007_alter_rustdesdevice_options_and_more'), + ] + + operations = [ + migrations.AddField( + model_name='rustdesdevice', + name='owner', + field=models.CharField(blank=True, max_length=100, verbose_name='所属用户'), + ), + ] diff --git a/api/models_user.py b/api/models_user.py index 65e9fff..6acd435 100644 --- a/api/models_user.py +++ b/api/models_user.py @@ -1,10 +1,12 @@ # cython:language_level=3 from django.db import models from django.contrib.auth.models import ( - BaseUserManager,AbstractBaseUser,PermissionsMixin + BaseUserManager, AbstractBaseUser, PermissionsMixin ) from .models_work import * from django.utils.translation import gettext as _ +from django.core.exceptions import ObjectDoesNotExist + class MyUserManager(BaseUserManager): def create_user(self, username, password=None): @@ -12,73 +14,130 @@ class MyUserManager(BaseUserManager): raise ValueError('Users must have an username') user = self.model(username=username, - ) - + ) + user.set_password(password) user.save(using=self._db) return user - + def create_superuser(self, username, password): user = self.create_user(username, - password=password, - - ) + password=password, + + ) user.is_admin = True user.save(using=self._db) return user class UserProfile(AbstractBaseUser, PermissionsMixin): - username = models.CharField(_('用户名'), + username = models.CharField(_('用户名'), unique=True, max_length=50) - + rid = models.CharField(verbose_name='RustDesk ID', max_length=16) uuid = models.CharField(verbose_name='uuid', max_length=60) autoLogin = models.BooleanField(verbose_name='autoLogin', default=True) rtype = models.CharField(verbose_name='rtype', max_length=20) deviceInfo = models.TextField(verbose_name=_('登录信息:'), blank=True) - + is_active = models.BooleanField(verbose_name=_('是否激活'), default=True) is_admin = models.BooleanField(verbose_name=_('是否管理员'), default=False) objects = MyUserManager() - + USERNAME_FIELD = 'username' # 用作用户名的字段 - REQUIRED_FIELDS = ['password'] #必须填写的字段 - - + REQUIRED_FIELDS = ['password'] # 必须填写的字段 + def get_full_name(self): # The user is identified by their email address return self.username - + def get_short_name(self): # The user is identified by their email address return self.username - + def __str__(self): # __unicode__ on Python 2 return self.username - - def has_perm(self, perm, obj=None): #有没有指定的权限 + + def has_perm(self, perm, obj=None): # 有没有指定的权限 "Does the user have a specific permission?" # Simplest possible answer: Yes, always return True - + def has_module_perms(self, app_label): "Does the user have permissions to view the app `app_label`?" # Simplest possible answer: Yes, always return True - + @classmethod + def get_uid_by_username(cls, username: str) -> int: + ''' + Get the uid by username. + + Args: + username (required): str, username + + Returns: + int: The uid of the user. + ''' + try: + user = cls.objects.get(username=username) + return user.pk + except ObjectDoesNotExist: + return None + + @classmethod + def get_username_by_uid(cls, uid: str) -> str: + ''' + Get the username by uid. + + Args: + uid (required): str, uid + + Returns: + str: The username of the user. + ''' + try: + user = cls.objects.get(pk=uid) + return user.username + except ObjectDoesNotExist: + return None + + @classmethod + def check_if_user_exists_by_username(cls, username: str) -> bool: + ''' + Check if a user exists. + + Args: + username (required): str, username + + Returns: + bool: True if the user exists, False otherwise. + ''' + return cls.objects.filter(username=username).exists() + + @classmethod + def check_if_user_exists_by_uid(cls, uid: str) -> bool: + ''' + Check if a user exists. + + Args: + uid (required): str, uid + + Returns: + bool: True if the user exists, False otherwise. + ''' + return cls.objects.filter(pk=uid).exists() @property def is_staff(self): "Is the user a member of staff?" - # Simplest possible answer: All admins are staff + # Simplest possible answer: All admins are staff return self.is_admin class Meta: - + verbose_name = _("用户") verbose_name_plural = _("用户列表") permissions = ( @@ -86,4 +145,3 @@ class UserProfile(AbstractBaseUser, PermissionsMixin): ("change_task_status", "Can change the status of tasks"), ("close_task", "Can remove a task by setting its status as closed"), ) - diff --git a/api/models_work.py b/api/models_work.py index 2991d83..7cdfe20 100644 --- a/api/models_work.py +++ b/api/models_work.py @@ -2,6 +2,9 @@ from django.db import models from django.contrib import admin from django.utils.translation import gettext as _ +from django.db.models import Q +from django.db.models.query import QuerySet + class RustDeskToken(models.Model): ''' Token @@ -12,17 +15,19 @@ class RustDeskToken(models.Model): uuid = models.CharField(verbose_name=_('uuid'), max_length=60) access_token = models.CharField(verbose_name=_('access_token'), max_length=60, blank=True) create_time = models.DateTimeField(verbose_name=_('登录时间'), auto_now_add=True) - #expire_time = models.DateTimeField(verbose_name='过期时间') + # expire_time = models.DateTimeField(verbose_name='过期时间') + class Meta: ordering = ('-username',) verbose_name = "Token" - verbose_name_plural = _("Token列表") + verbose_name_plural = _("Token列表") + class RustDeskTokenAdmin(admin.ModelAdmin): list_display = ('username', 'uid') search_fields = ('username', 'uid') - list_filter = ('create_time', ) #过滤器 - + list_filter = ('create_time', ) # 过滤器 + class RustDeskTag(models.Model): ''' Tags @@ -30,17 +35,18 @@ class RustDeskTag(models.Model): uid = models.CharField(verbose_name=_('所属用户ID'), max_length=16) tag_name = models.CharField(verbose_name=_('标签名称'), max_length=60) tag_color = models.CharField(verbose_name=_('标签颜色'), max_length=60, blank=True) - + class Meta: ordering = ('-uid',) verbose_name = "Tags" verbose_name_plural = _("Tags列表") + class RustDeskTagAdmin(admin.ModelAdmin): list_display = ('tag_name', 'uid', 'tag_color') search_fields = ('tag_name', 'uid') list_filter = ('uid', ) - + class RustDeskPeer(models.Model): ''' Pees @@ -53,19 +59,32 @@ class RustDeskPeer(models.Model): platform = models.CharField(verbose_name=_('平台'), max_length=30) tags = models.CharField(verbose_name=_('标签'), max_length=30) rhash = models.CharField(verbose_name=_('设备链接密码'), max_length=60) - + class Meta: ordering = ('-username',) verbose_name = "Peers" - verbose_name_plural = _("Peers列表" ) - + verbose_name_plural = _("Peers列表") + + @classmethod + def get_peers_by_uid(cls, uid: str) -> QuerySet: + ''' + Get all peers owned by a user. + + Args: + uid (required): str, uid + + Returns: + QuerySet: A queryset of peers owned by the user. + ''' + return cls.objects.filter(uid=uid) + class RustDeskPeerAdmin(admin.ModelAdmin): list_display = ('rid', 'uid', 'username', 'hostname', 'platform', 'alias', 'tags') search_fields = ('deviceid', 'alias') list_filter = ('rid', 'uid', ) - - + + class RustDesDevice(models.Model): rid = models.CharField(verbose_name=_('客户端ID'), max_length=60, blank=True) cpu = models.CharField(verbose_name='CPU', max_length=100) @@ -77,19 +96,83 @@ class RustDesDevice(models.Model): version = models.CharField(verbose_name=_('客户端版本'), max_length=100) create_time = models.DateTimeField(verbose_name=_('设备注册时间'), auto_now_add=True) update_time = models.DateTimeField(verbose_name=('设备更新时间'), auto_now=True, blank=True) - + owner = models.CharField(verbose_name=_('所属用户'), max_length=100, blank=True) + class Meta: ordering = ('-rid',) verbose_name = _("设备") - verbose_name_plural = _("设备列表" ) - + verbose_name_plural = _("设备列表") + + @classmethod + def update_device(cls, rid: str, uuid: str, cpu: str = None, + hostname: str = None, memory: str = None, os: str = None, + username: str = None, version: str = None, owner: str = None + ) -> 'RustDesDevice': + ''' + Update or create a device and update its information if provided. + + Args: + rid (required): str, rid + uuid (required): str, uuid + cpu (optional): str, default None + hostname (optional): str, default None + memory (optional): str, default None + os (optional): str, default None + username (optional): str, default None + version (optional): str, default None + owner (optional): str, default None + + Returns: + RustDesDevice: The created or updated device object. + ''' + # Use get_or_create to either get an existing device or create a new one + device, created = cls.objects.get_or_create(rid=rid, uuid=uuid) + # Collect all non-None fields that need to be updated + update_fields = {} + + if cpu is not None: + update_fields['cpu'] = cpu + if hostname is not None: + update_fields['hostname'] = hostname + if memory is not None: + update_fields['memory'] = memory + if os is not None: + update_fields['os'] = os + if username is not None: + update_fields['username'] = username + if version is not None: + update_fields['version'] = version + if owner is not None: # assuming you want to update owner even when it is None + update_fields['owner'] = owner + # Update only the provided fields + if update_fields: + for field, value in update_fields.items(): + setattr(device, field, value) + device.save(update_fields=update_fields.keys()) + return device + + @classmethod + def get_devices_by_owner(cls, owner: str) -> QuerySet: + ''' + Get all devices owned by a user. + + Args: + owner (required): str, owner + + Returns: + QuerySet: A queryset of devices owned by the user. + ''' + return cls.objects.filter(owner=owner) + + class RustDesDeviceAdmin(admin.ModelAdmin): list_display = ('rid', 'hostname', 'memory', 'uuid', 'version', 'create_time', 'update_time') search_fields = ('hostname', 'memory') list_filter = ('rid', ) + class ConnLog(models.Model): - id = models.IntegerField(verbose_name='ID',primary_key=True) + id = models.IntegerField(verbose_name='ID', primary_key=True) action = models.CharField(verbose_name='Action', max_length=20, null=True) conn_id = models.CharField(verbose_name='Connection ID', max_length=10, null=True) from_ip = models.CharField(verbose_name='From IP', max_length=30, null=True) @@ -100,13 +183,15 @@ class ConnLog(models.Model): session_id = models.CharField(verbose_name='Session ID', max_length=60, null=True) uuid = models.CharField(verbose_name='uuid', max_length=60, null=True) + class ConnLogAdmin(admin.ModelAdmin): list_display = ('id', 'action', 'conn_id', 'from_ip', 'from_id', 'rid', 'conn_start', 'conn_end', 'session_id', 'uuid') search_fields = ('from_ip', 'rid') list_filter = ('id', 'from_ip', 'from_id', 'rid', 'conn_start', 'conn_end') + class FileLog(models.Model): - id = models.IntegerField(verbose_name='ID',primary_key=True) + id = models.IntegerField(verbose_name='ID', primary_key=True) file = models.CharField(verbose_name='Path', max_length=500) remote_id = models.CharField(verbose_name='Remote ID', max_length=20, default='0') user_id = models.CharField(verbose_name='User ID', max_length=20, default='0') @@ -115,11 +200,13 @@ class FileLog(models.Model): direction = models.IntegerField(verbose_name='Direction', default=0) logged_at = models.DateTimeField(verbose_name='Logged At', null=True) + class FileLogAdmin(admin.ModelAdmin): list_display = ('id', 'file', 'remote_id', 'user_id', 'user_ip', 'filesize', 'direction', 'logged_at') search_fields = ('file', 'remote_id', 'user_id', 'user_ip') list_filter = ('id', 'file', 'remote_id', 'user_id', 'user_ip', 'filesize', 'direction', 'logged_at') + class ShareLink(models.Model): ''' 分享链接 ''' @@ -129,16 +216,14 @@ class ShareLink(models.Model): is_used = models.BooleanField(verbose_name=_('是否使用'), default=False) is_expired = models.BooleanField(verbose_name=_('是否过期'), default=False) create_time = models.DateTimeField(verbose_name=_('生成时间'), auto_now_add=True) - - class Meta: ordering = ('-create_time',) verbose_name = _("分享链接") - verbose_name_plural = _("链接列表" ) - + verbose_name_plural = _("链接列表") + class ShareLinkAdmin(admin.ModelAdmin): list_display = ('shash', 'uid', 'peers', 'is_used', 'is_expired', 'create_time') search_fields = ('peers', ) - list_filter = ('is_used', 'uid', 'is_expired' ) \ No newline at end of file + list_filter = ('is_used', 'uid', 'is_expired') diff --git a/api/views_api.py b/api/views_api.py index 83ca601..f537706 100644 --- a/api/views_api.py +++ b/api/views_api.py @@ -12,6 +12,97 @@ from django.db.models import Q import copy from .views_front import * from django.utils.translation import gettext as _ +from django.conf import settings + + +def _get_uid_by_username(username: str | int) -> int: + """ + Get user ID by username. + args: + username: The username to be queried. + """ + try: + uid = int(username) + if not UserProfile.check_if_user_exists_by_uid(uid): + uid = None + except: + uid = UserProfile.get_uid_by_username(username=username) + return uid + + +def update_device_owner(rid: str, uuid: str, owner: str = None, add: bool = False, *args, **kwargs) -> RustDesDevice: + """ + Update device owner. + + Args: + rid: The device ID to be updated. + uuid: The device UUID to be updated. + owner: The owner to be set. If add is False, owner will be set to None. + add: If True, add owner; if False, remove owner by setting it to None. + """ + # If add is False, set owner to None, indicating the owner will be removed + if not add: + owner = "" + # Retrieve the device using both 'rid' and 'uuid' for filtering + device = RustDesDevice.update_device(rid=rid, uuid=uuid, owner=owner) + return device + + +def update_self_tag(username: str) -> RustDeskTag: + """ + Update self tag. + """ + uid = _get_uid_by_username(username=username) + if not uid: + return None + tag, created = RustDeskTag.objects.get_or_create(uid=uid, tag_name=settings.SELF_TAG_NAME) + tag.tag_color = settings.SELF_TAG_COLOR + tag.save() + return tag + + +def update_self_devices_to_peers(username: str) -> None: + """ + Update self devices to peers. + """ + self_tag = update_self_tag(username=username) + if self_tag is None: + return None + username = UserProfile.get_username_by_uid(uid=self_tag.uid) + devices = RustDesDevice.get_devices_by_owner(owner=username) + for device in devices: + device: RustDesDevice = device + peer, created = RustDeskPeer.objects.get_or_create(uid=self_tag.uid, rid=device.rid) + peer.username = device.username + peer.hostname = device.hostname + peer.platform = device.os + tags = peer.tags.split(',') + if self_tag.tag_name not in tags: + tags.append(self_tag.tag_name) + peer.tags = ",".join(tags) + peer.save() + + +def delete_self_devices_from_peers(username: str) -> None: + """ + When a user logs out, delete the devices from peers. + """ + uid = _get_uid_by_username(username=username) + if not uid: + return None + peers = RustDeskPeer.get_peers_by_uid(uid=uid) + for peer in peers: + peer: RustDeskPeer = peer + tags = peer.tags.split(',') + if settings.SELF_TAG_NAME in tags: + l = len(tags) + if l == 1 or l == 2: + peer.delete() + else: + tags.remove(settings.SELF_TAG_NAME) + peer.tags = ",".join(tags) + peer.rhash = "" + peer.save() def login(request): @@ -21,7 +112,6 @@ def login(request): return JsonResponse(result) data = json.loads(request.body.decode()) - username = data.get('username', '') password = data.get('password', '') rid = data.get('id', '') @@ -39,7 +129,8 @@ def login(request): user.rtype = rtype user.deviceInfo = json.dumps(deviceInfo) user.save() - + update_device_owner(rid=rid, uuid=uuid, owner=username, add=True, **deviceInfo) + update_self_tag(username=username) token = RustDeskToken.objects.filter(Q(uid=user.id) & Q(username=user.username) & Q(rid=user.rid)).first() # 检查是否过期 @@ -84,6 +175,8 @@ def logout(request): token.delete() result = {'code': 1} + update_device_owner(rid=rid, uuid=uuid, owner=None, add=False) + delete_self_devices_from_peers(username=user.username) return JsonResponse(result) @@ -120,8 +213,8 @@ def ab(request): if not token: result = {'error': _('拉取列表错误！')} return JsonResponse(result) - if request.method == 'GET': + update_self_devices_to_peers(username=token.uid) result = {} uid = token.uid tags = RustDeskTag.objects.filter(Q(uid=uid)) @@ -180,6 +273,9 @@ def ab(request): RustDeskPeer.objects.filter(uid=token.uid).delete() newlist = [] for one in peers: + tags = one['tags'] + if settings.SELF_TAG_NAME in tags: + tags.remove(settings.SELF_TAG_NAME) peer = RustDeskPeer( uid=token.uid, rid=one['id'], @@ -187,13 +283,13 @@ def ab(request): hostname=one['hostname'], alias=one['alias'], platform=one['platform'], - tags=','.join(one['tags']), + tags=','.join(tags), rhash=one['hash'], - - ) newlist.append(peer) RustDeskPeer.objects.bulk_create(newlist) + # update self devices to peers after updating the address book + update_self_devices_to_peers(username=token.uid) result = { 'code': 102, @@ -208,14 +304,14 @@ def ab_get(request): return ab(request) -def sysinfo(request): +def sysinfo(request) -> JsonResponse: # 客户端注册服务后，才会发送设备信息 result = {} if request.method == 'GET': result['error'] = _('错误的提交方式！') return JsonResponse(result) - postdata = json.loads(request.body) + postdata: dict[str, str] = json.loads(request.body) device = RustDesDevice.objects.filter(Q(rid=postdata['id']) & Q(uuid=postdata['uuid'])).first() if not device: device = RustDesDevice( @@ -270,7 +366,7 @@ def audit(request): elif audit_type == "close": ConnLog.objects.filter(Q(conn_id=postdata['conn_id'])).update(conn_end=datetime.datetime.now()) elif 'is_file' in postdata: - print(postdata) + # print(postdata) files = json.loads(postdata['info'])['files'] filesize = convert_filesize(int(files[0][1])) new_file_log = FileLog( diff --git a/api/views_front.py b/api/views_front.py index fa43427..cf265eb 100644 --- a/api/views_front.py +++ b/api/views_front.py @@ -228,7 +228,6 @@ def get_all_info(): device = devices.get(peer.rid, None) if device: devices[peer.rid]['rust_user'] = user.username - for k, v in devices.items(): devices[k]['status'] = _('在线') if (now - datetime.datetime.strptime(v['update_time'], '%Y-%m-%d %H:%M')).seconds <= 120 else _('离线') return [v for k, v in devices.items()] diff --git a/db/db.sqlite3 b/db/db.sqlite3 index afb6fe0444c5438fbb5e8504d213d663f8555df5..29ede4e26e873fd0dd8710c0385faec747caa28d 100644 GIT binary patch delta 352 zcmZo@;A&XFH9=ZXlYxN&gq4Aq5r~5(>KL}eUBTEGX11lo~ zD+6OaOG`_0L*vc9k|)?012-!wSo0f6@-tg<6JT@jv4Spv%uON;YR&vx_&U00~P=#i8~y~UDAg%#|>=@%@R#2GcW-!W$DWn|*u Z+@5X9G>?e|7=HV<^I9-{c3;#W000xgWxW6Z delta 175 zcmZozz}3*eH9=ZXgMonogkivcqK+}E27_MFf{iH)_?aa+eJ8UE7;jco5a!(MD|v#A zF>SM8LIA(4BpR5X-P(WL26NMW^r+5exA7q-{fX}KSs&sOZwX{=`-qm zSkxde&47tjn)wa`-%Z|+ygb~jTv43IIXKxiuokiuv+6M4+3YB=hk5%QW2Rn4CI*J> V*``eMm{^*vP1(2eS}=Wf2LPPrE+GH_ diff --git a/rustdesk_server_api/settings.py b/rustdesk_server_api/settings.py index 639e615..e094cf0 100644 --- a/rustdesk_server_api/settings.py +++ b/rustdesk_server_api/settings.py @@ -172,3 +172,7 @@ LANGUAGES = ( LOCALE_PATHS = ( os.path.join(BASE_DIR, 'locale'), ) + + +SELF_TAG_NAME='self' +SELF_TAG_COLOR=0xFF6A1B9A \ No newline at end of file