users' devices will add to peers automatically

This commit is contained in:
Tao Chen 2024-08-29 00:21:47 +08:00
parent e2bbf9c04c
commit c44206cde6
8 changed files with 318 additions and 55 deletions

5
.gitignore vendored
View File

@ -28,4 +28,7 @@ LICENSE.rst
db/test_db.sqlite3
job2en.py
新建文本文档.txt
新建文本文档.txt
venv
.env

View File

@ -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='所属用户'),
),
]

View File

@ -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"),
)

View File

@ -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' )
list_filter = ('is_used', 'uid', 'is_expired')

View File

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

View File

@ -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()]

Binary file not shown.

View File

@ -172,3 +172,7 @@ LANGUAGES = (
LOCALE_PATHS = (
os.path.join(BASE_DIR, 'locale'),
)
SELF_TAG_NAME='self'
SELF_TAG_COLOR=0xFF6A1B9A