diff --git a/apps/authentication/models/connection_token.py b/apps/authentication/models/connection_token.py index 63049170b..eef633d95 100644 --- a/apps/authentication/models/connection_token.py +++ b/apps/authentication/models/connection_token.py @@ -172,7 +172,7 @@ class ConnectionToken(JMSOrgBaseModel): if not applet: return None - host_account = applet.select_host_account() + host_account = applet.select_host_account(self.user) if not host_account: raise JMSException({'error': 'No host account available'}) diff --git a/apps/terminal/migrations/0050_auto_20220606_1745.py b/apps/terminal/migrations/0050_auto_20220606_1745.py index e88d37971..d0eb6ea5d 100644 --- a/apps/terminal/migrations/0050_auto_20220606_1745.py +++ b/apps/terminal/migrations/0050_auto_20220606_1745.py @@ -4,7 +4,6 @@ from django.db import migrations, models class Migration(migrations.Migration): - dependencies = [ ('terminal', '0049_endpoint_redis_port'), ] @@ -13,10 +12,10 @@ class Migration(migrations.Migration): migrations.AlterField( model_name='terminal', name='type', - field=models.CharField(choices=[ - ('koko', 'KoKo'), ('guacamole', 'Guacamole'), ('omnidb', 'OmniDB'), - ('xrdp', 'Xrdp'), ('lion', 'Lion'), ('core', 'Core'), ('celery', 'Celery'), - ('magnus', 'Magnus'), ('razor', 'Razor'), ('tinker', 'Tinker'), - ], default='koko', max_length=64, verbose_name='type'), + field=models.CharField( + choices=[('koko', 'KoKo'), ('guacamole', 'Guacamole'), ('omnidb', 'OmniDB'), ('xrdp', 'Xrdp'), + ('lion', 'Lion'), ('core', 'Core'), ('celery', 'Celery'), ('magnus', 'Magnus'), + ('razor', 'Razor'), ('tinker', 'Tinker'), ('video_worker', 'Video Worker')], default='koko', + max_length=64, verbose_name='type'), ), ] diff --git a/apps/terminal/migrations/0061_applet_can_concurrent.py b/apps/terminal/migrations/0061_applet_can_concurrent.py new file mode 100644 index 000000000..4ca762e65 --- /dev/null +++ b/apps/terminal/migrations/0061_applet_can_concurrent.py @@ -0,0 +1,18 @@ +# Generated by Django 3.2.17 on 2023-05-09 11:02 + +from django.db import migrations, models + + +class Migration(migrations.Migration): + + dependencies = [ + ('terminal', '0060_sessionsharing_action_permission'), + ] + + operations = [ + migrations.AddField( + model_name='applet', + name='can_concurrent', + field=models.BooleanField(default=True, verbose_name='Can concurrent'), + ), + ] diff --git a/apps/terminal/models/applet/applet.py b/apps/terminal/models/applet/applet.py index c7296eda6..8bf6b3a85 100644 --- a/apps/terminal/models/applet/applet.py +++ b/apps/terminal/models/applet/applet.py @@ -32,6 +32,7 @@ class Applet(JMSBaseModel): is_active = models.BooleanField(default=True, verbose_name=_('Is active')) builtin = models.BooleanField(default=False, verbose_name=_('Builtin')) protocols = models.JSONField(default=list, verbose_name=_('Protocol')) + can_concurrent = models.BooleanField(default=True, verbose_name=_('Can concurrent')) tags = models.JSONField(default=list, verbose_name=_('Tags')) comment = models.TextField(default='', blank=True, verbose_name=_('Comment')) hosts = models.ManyToManyField( @@ -134,37 +135,68 @@ class Applet(JMSBaseModel): shutil.copytree(path, pkg_path) return instance, serializer - def select_host_account(self): - # 选择激活的发布机 + def select_host(self, user): hosts = [ host for host in self.hosts.filter(is_active=True) if host.load != 'offline' ] - if not hosts: return None - key_tmpl = 'applet_host_accounts_{}_{}' - host = random.choice(hosts) - using_keys = cache.keys(key_tmpl.format(host.id, '*')) or [] - accounts_username_used = list(cache.get_many(using_keys).values()) - logger.debug('Applet host account using: {}: {}'.format(host.name, accounts_username_used)) - accounts = host.accounts.all() \ - .filter(is_active=True, privileged=False) \ - .exclude(username__in=accounts_username_used) + prefer_key = 'applet_host_prefer_{}'.format(user.id) + prefer_host_id = cache.get(prefer_key, None) + pref_host = [host for host in hosts if host.id == prefer_host_id] + if pref_host: + host = pref_host[0] + else: + host = random.choice(hosts) + cache.set(prefer_key, host.id, timeout=None) + return host - msg = 'Applet host remain accounts: {}: {}'.format(host.name, len(accounts)) + @staticmethod + def random_select_prefer_account(user, host, accounts): + msg = 'Applet host remain public accounts: {}: {}'.format(host.name, len(accounts)) if len(accounts) == 0: logger.error(msg) - else: - logger.debug(msg) - - if not accounts: return None + prefer_host_account_key = 'applet_host_prefer_account_{}_{}'.format(user.id, host.id) + prefer_account_id = cache.get(prefer_host_account_key, None) + prefer_account = accounts.filter(id=prefer_account_id).first() + if prefer_account: + account = prefer_account + else: + account = random.choice(accounts) + cache.set(prefer_host_account_key, account.id, timeout=None) + return account - account = random.choice(accounts) + def select_host_account(self, user): + # 选择激活的发布机 + host = self.select_host(user) + if not host: + return None + can_concurrent = self.can_concurrent and self.type == 'general' + + accounts = host.accounts.all().filter(is_active=True, privileged=False) + private_account = accounts.filter(username='js_{}'.format(user.username)).first() + accounts_using_key_tmpl = 'applet_host_accounts_{}_{}' + + if private_account and can_concurrent: + account = private_account + else: + using_keys = cache.keys(accounts_using_key_tmpl.format(host.id, '*')) or [] + accounts_username_used = list(cache.get_many(using_keys).values()) + logger.debug('Applet host account using: {}: {}'.format(host.name, accounts_username_used)) + + # 优先使用 private account + if private_account and private_account.username not in accounts_username_used: + account = private_account + else: + accounts = accounts.exclude(username__in=accounts_username_used) + account = self.random_select_prefer_account(user, host, accounts) + if not account: + return ttl = 60 * 60 * 24 - lock_key = key_tmpl.format(host.id, account.username) + lock_key = accounts_using_key_tmpl.format(host.id, account.username) cache.set(lock_key, account.username, ttl) return { diff --git a/apps/terminal/models/applet/host.py b/apps/terminal/models/applet/host.py index 47de718df..1ad129baa 100644 --- a/apps/terminal/models/applet/host.py +++ b/apps/terminal/models/applet/host.py @@ -84,9 +84,13 @@ class AppletHost(Host): return random_string(16, special_char=True) def generate_accounts(self): - amount = int(os.getenv('TERMINAL_ACCOUNTS_AMOUNT', 100)) - now_count = self.accounts.filter(privileged=False).count() - need = amount - now_count + self.generate_public_accounts() + self.generate_private_accounts() + + def generate_public_accounts(self): + public_amount = int(os.getenv('TERMINAL_ACCOUNTS_AMOUNT', 100)) + now_count = self.accounts.filter(privileged=False, username__startswith='jms').count() + need = public_amount - now_count accounts = [] account_model = self.accounts.model @@ -99,7 +103,30 @@ class AppletHost(Host): org_id=self.LOCKING_ORG, is_active=False, ) accounts.append(account) - bulk_create_with_history(accounts, account_model, batch_size=20) + bulk_create_with_history(accounts, account_model, batch_size=20, ignore_conflicts=True) + + def generate_private_accounts_by_usernames(self, usernames): + accounts = [] + account_model = self.accounts.model + for username in usernames: + password = self.random_password() + username = 'js_' + username + account = account_model( + username=username, secret=password, name=username, + asset_id=self.id, secret_type='password', version=1, + org_id=self.LOCKING_ORG, is_active=False, + ) + accounts.append(account) + bulk_create_with_history(accounts, account_model, batch_size=20, ignore_conflicts=True) + + def generate_private_accounts(self): + from users.models import User + usernames = User.objects \ + .filter(is_active=True, is_service_account=False) \ + .values_list('username', flat=True) + account_usernames = self.accounts.all().values_list('username', flat=True) + not_exist_users = set(usernames) - set(account_usernames) + self.generate_private_accounts_by_usernames(not_exist_users) class AppletHostDeployment(JMSBaseModel): diff --git a/apps/terminal/signal_handlers/applet.py b/apps/terminal/signal_handlers/applet.py index 18595fc7c..4fe390b8e 100644 --- a/apps/terminal/signal_handlers/applet.py +++ b/apps/terminal/signal_handlers/applet.py @@ -2,11 +2,14 @@ from django.db.models.signals import post_save, post_delete from django.dispatch import receiver from django.utils.functional import LazyObject +from accounts.models import Account from common.signals import django_ready from common.utils import get_logger from common.utils.connection import RedisPubSub from orgs.utils import tmp_to_builtin_org +from users.models import User from ..models import Applet, AppletHost +from ..tasks import applet_host_generate_accounts from ..utils import DBPortManager db_port_manager: DBPortManager @@ -19,12 +22,30 @@ def on_applet_host_create(sender, instance, created=False, **kwargs): return applets = Applet.objects.all() instance.applets.set(applets) - with tmp_to_builtin_org(system=1): - instance.generate_accounts() + applet_host_generate_accounts.delay(instance.id) applet_host_change_pub_sub.publish(True) +@receiver(post_save, sender=User) +def on_user_create_create_account(sender, instance, created=False, **kwargs): + if not created: + return + + with tmp_to_builtin_org(system=1): + applet_hosts = AppletHost.objects.all() + for host in applet_hosts: + host.generate_private_accounts_by_usernames([instance.username]) + + +@receiver(post_delete, sender=User) +def on_user_delete_remove_account(sender, instance, **kwargs): + with tmp_to_builtin_org(system=1): + applet_hosts = AppletHost.objects.all().values_list('id', flat=True) + accounts = Account.objects.filter(asset_id__in=applet_hosts, username=instance.username) + accounts.delete() + + @receiver(post_delete, sender=AppletHost) def on_applet_host_delete(sender, instance, **kwargs): applet_host_change_pub_sub.publish(True) diff --git a/apps/terminal/tasks.py b/apps/terminal/tasks.py index 28356972f..66e6871d7 100644 --- a/apps/terminal/tasks.py +++ b/apps/terminal/tasks.py @@ -16,7 +16,7 @@ from ops.celery.decorator import ( from orgs.utils import tmp_to_builtin_org from .backends import server_replay_storage from .models import ( - Status, Session, Task, AppletHostDeployment + Status, Session, Task, AppletHostDeployment, AppletHost ) from .utils import find_session_replay_local @@ -82,7 +82,7 @@ def upload_session_replay_to_external_storage(session_id): @shared_task( verbose_name=_('Run applet host deployment'), - activity_callback=lambda self, did, *args, **kwargs: ([did], ) + activity_callback=lambda self, did, *args, **kwargs: ([did],) ) def run_applet_host_deployment(did): with tmp_to_builtin_org(system=1): @@ -98,3 +98,16 @@ def run_applet_host_deployment_install_applet(did, applet_id): with tmp_to_builtin_org(system=1): deployment = AppletHostDeployment.objects.get(id=did) deployment.install_applet(applet_id) + + +@shared_task( + verbose_name=_('Generate applet host accounts'), + activity_callback=lambda self, host_id, *args, **kwargs: ([host_id],) +) +def applet_host_generate_accounts(host_id): + applet_host = AppletHost.objects.filter(id=host_id).first() + if not applet_host: + return + + with tmp_to_builtin_org(system=1): + applet_host.generate_accounts()