diff --git a/apps/authentication/api/connection_token.py b/apps/authentication/api/connection_token.py index 41cce13f7..18cc0d5cb 100644 --- a/apps/authentication/api/connection_token.py +++ b/apps/authentication/api/connection_token.py @@ -14,16 +14,16 @@ from rest_framework.request import Request from rest_framework.response import Response from rest_framework.serializers import ValidationError +from assets.const import CloudTypes from common.api import JMSModelViewSet from common.exceptions import JMSException -from common.utils import random_string +from common.utils import random_string, get_logger from common.utils.django import get_request_os from common.utils.http import is_true from orgs.mixins.api import RootOrgViewMixin from perms.models import ActionChoices from terminal.connect_methods import NativeClient, ConnectMethodUtil from terminal.models import EndpointRule -from assets.const import CloudTypes from ..models import ConnectionToken from ..serializers import ( ConnectionTokenSerializer, ConnectionTokenSecretSerializer, @@ -31,6 +31,7 @@ from ..serializers import ( ) __all__ = ['ConnectionTokenViewSet', 'SuperConnectionTokenViewSet'] +logger = get_logger(__name__) class RDPFileClientProtocolURLMixin: @@ -363,5 +364,11 @@ class SuperConnectionTokenViewSet(ConnectionTokenViewSet): @action(methods=['DELETE', 'POST'], detail=False, url_path='applet-account/release') def release_applet_account(self, *args, **kwargs): account_id = self.request.data.get('id') - msg = ConnectionToken.release_applet_account(account_id) - return Response({'msg': msg}) + released = ConnectionToken.release_applet_account(account_id) + + if released: + logger.debug('Release applet account success: {}'.format(account_id)) + return Response({'msg': 'released'}) + else: + logger.error('Release applet account error: {}'.format(account_id)) + return Response({'error': 'not found or expired'}, status=400) diff --git a/apps/authentication/models/connection_token.py b/apps/authentication/models/connection_token.py index 385294122..958ab8fa9 100644 --- a/apps/authentication/models/connection_token.py +++ b/apps/authentication/models/connection_token.py @@ -11,7 +11,8 @@ from rest_framework.exceptions import PermissionDenied from assets.const import Protocol from common.db.fields import EncryptCharField -from common.utils import lazyproperty, pretty_string, bulk_get, reverse +from common.exceptions import JMSException +from common.utils import lazyproperty, pretty_string, bulk_get from common.utils.timezone import as_current_tz from orgs.mixins.models import JMSOrgBaseModel from terminal.models import Applet @@ -172,7 +173,7 @@ class ConnectionToken(JMSOrgBaseModel): host_account = applet.select_host_account() if not host_account: - return None + raise JMSException({'error': 'No host account available'}) host, account, lock_key, ttl = bulk_get(host_account, ('host', 'account', 'lock_key', 'ttl')) gateway = host.gateway.select_gateway() if host.domain else None @@ -196,8 +197,7 @@ class ConnectionToken(JMSOrgBaseModel): if lock_key: cache.delete(lock_key) cache.delete(token_account_relate_key) - return 'released' - return 'not found or expired' + return True @lazyproperty def account_object(self): diff --git a/apps/orgs/models.py b/apps/orgs/models.py index 2fda6b443..de2b087a3 100644 --- a/apps/orgs/models.py +++ b/apps/orgs/models.py @@ -134,12 +134,12 @@ class Organization(OrgRoleMixin, JMSBaseModel): return self.id @classmethod - def get_or_create_builtin(cls, name, **kwargs): - _id = kwargs.get('id') - org = cls.get_instance(cls.DEFAULT_ID) + def get_or_create_builtin(cls, **kwargs): + _id = kwargs['id'] + org = cls.get_instance(_id) if org: return org - org, created = cls.objects.get_or_create(name=name, defaults=kwargs) + org, created = cls.objects.get_or_create(id=_id, defaults=kwargs) if created: org.builtin = True org.save() diff --git a/apps/terminal/models/applet/applet.py b/apps/terminal/models/applet/applet.py index da87eb12f..1fd8a670f 100644 --- a/apps/terminal/models/applet/applet.py +++ b/apps/terminal/models/applet/applet.py @@ -11,7 +11,9 @@ from django.utils.translation import gettext_lazy as _ from rest_framework.serializers import ValidationError from common.db.models import JMSBaseModel -from common.utils import lazyproperty +from common.utils import lazyproperty, get_logger + +logger = get_logger(__name__) __all__ = ['Applet', 'AppletPublication'] @@ -109,19 +111,27 @@ class Applet(JMSBaseModel): if not hosts: return None + key_tmpl = 'applet_host_accounts_{}_{}' host = random.choice(hosts) - using_keys = cache.keys('host_accounts_{}_*'.format(host.id)) or [] - accounts_used = cache.get_many(using_keys) - accounts = host.accounts.all().exclude(username__in=accounts_used) + 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) + + msg = 'Applet host remain accounts: {}: {}'.format(host.name, len(accounts)) + if len(accounts) == 0: + logger.error(msg) + else: + logger.debug(msg) - if not accounts: - accounts = host.accounts.all() if not accounts: return None account = random.choice(accounts) ttl = 60 * 60 * 24 - lock_key = 'applet_host_accounts_{}_{}'.format(host.id, account.username) + lock_key = key_tmpl.format(host.id, account.username) cache.set(lock_key, account.username, ttl) return { @@ -131,11 +141,6 @@ class Applet(JMSBaseModel): 'ttl': ttl } - @staticmethod - def release_host_and_account(host_id, username): - key = 'applet_host_accounts_{}_{}'.format(host_id, username) - cache.delete(key) - class AppletPublication(JMSBaseModel): applet = models.ForeignKey('Applet', on_delete=models.CASCADE, related_name='publications',