From 3abc8bddfa30cce09f0ca653591ea4c6557b5bf7 Mon Sep 17 00:00:00 2001
From: feng <1304903146@qq.com>
Date: Wed, 13 Sep 2023 16:52:11 +0800
Subject: [PATCH] =?UTF-8?q?feat:=20=E7=94=A8=E6=88=B7=E5=9C=A8=E7=BA=BFses?=
 =?UTF-8?q?sion=E6=8E=A7=E5=88=B6?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 apps/audits/api.py                        |  5 ++-
 apps/audits/signal_handlers/login_log.py  | 35 +++++++++++++++-
 apps/audits/utils.py                      |  8 ++--
 apps/rbac/builtin.py                      |  1 +
 apps/rbac/const.py                        |  1 +
 apps/rbac/tree.py                         |  1 +
 apps/terminal/api/session/session.py      |  2 +-
 apps/terminal/models/session/__init__.py  |  2 +-
 apps/terminal/signal_handlers/session.py  |  1 +
 apps/users/api/__init__.py                |  1 +
 apps/users/api/session.py                 | 51 +++++++++++++++++++++++
 apps/users/migrations/0044_usersession.py | 36 ++++++++++++++++
 apps/users/models/__init__.py             |  1 +
 apps/users/models/session.py              | 40 ++++++++++++++++++
 apps/users/serializers/__init__.py        |  1 +
 apps/users/serializers/session.py         | 32 ++++++++++++++
 apps/users/signal_handlers.py             | 18 +++++++-
 apps/users/urls/api_urls.py               |  1 +
 18 files changed, 227 insertions(+), 10 deletions(-)
 create mode 100644 apps/users/api/session.py
 create mode 100644 apps/users/migrations/0044_usersession.py
 create mode 100644 apps/users/models/session.py
 create mode 100644 apps/users/serializers/session.py

diff --git a/apps/audits/api.py b/apps/audits/api.py
index c635e4170..490608e22 100644
--- a/apps/audits/api.py
+++ b/apps/audits/api.py
@@ -26,7 +26,10 @@ from terminal.models import default_storage
 from users.models import User
 from .backends import TYPE_ENGINE_MAPPING
 from .const import ActivityChoices
-from .models import FTPLog, UserLoginLog, OperateLog, PasswordChangeLog, ActivityLog, JobLog
+from .models import (
+    FTPLog, UserLoginLog, OperateLog, PasswordChangeLog,
+    ActivityLog, JobLog,
+)
 from .serializers import (
     FTPLogSerializer, UserLoginLogSerializer, JobLogSerializer,
     OperateLogSerializer, OperateLogActionDetailSerializer,
diff --git a/apps/audits/signal_handlers/login_log.py b/apps/audits/signal_handlers/login_log.py
index 8a4e5d262..148196699 100644
--- a/apps/audits/signal_handlers/login_log.py
+++ b/apps/audits/signal_handlers/login_log.py
@@ -1,17 +1,23 @@
 # -*- coding: utf-8 -*-
 #
+from datetime import timedelta
+from importlib import import_module
+
 from django.conf import settings
 from django.contrib.auth import BACKEND_SESSION_KEY
+from django.core.cache import caches
 from django.dispatch import receiver
 from django.utils import timezone, translation
 from django.utils.functional import LazyObject
 from django.utils.translation import gettext_lazy as _
 from rest_framework.request import Request
 
+from audits.models import UserLoginLog
 from authentication.signals import post_auth_failed, post_auth_success
 from authentication.utils import check_different_city_login_if_need
 from common.utils import get_request_ip, get_logger
-from users.models import User
+from users.models import User, UserSession
+from ..const import LoginTypeChoices
 from ..utils import write_login_log
 
 logger = get_logger(__name__)
@@ -75,6 +81,26 @@ def generate_data(username, request, login_type=None):
     return data
 
 
+def create_user_session(session_key, user_id, instance: UserLoginLog):
+    session_store_cls = import_module(settings.SESSION_ENGINE).SessionStore
+    session_store = session_store_cls(session_key=session_key)
+    cache_key = session_store.cache_key
+    ttl = caches[settings.SESSION_CACHE_ALIAS].ttl(cache_key)
+
+    online_session_data = {
+        'user_id': user_id,
+        'ip': instance.ip,
+        'key': session_key,
+        'city': instance.city,
+        'type': instance.type,
+        'backend': instance.backend,
+        'user_agent': instance.user_agent,
+        'date_created': instance.datetime,
+        'date_expired': instance.datetime + timedelta(seconds=ttl),
+    }
+    UserSession.objects.create(**online_session_data)
+
+
 @receiver(post_auth_success)
 def on_user_auth_success(sender, user, request, login_type=None, **kwargs):
     logger.debug('User login success: {}'.format(user.username))
@@ -84,7 +110,12 @@ def on_user_auth_success(sender, user, request, login_type=None, **kwargs):
     )
     request.session['login_time'] = data['datetime'].strftime("%Y-%m-%d %H:%M:%S")
     data.update({'mfa': int(user.mfa_enabled), 'status': True})
-    write_login_log(**data)
+    instance = write_login_log(**data)
+    session_key = request.session.session_key
+    # TODO 目前只记录 web 登录的 session
+    if not session_key or instance.type != LoginTypeChoices.web:
+        return
+    create_user_session(session_key, user.id, instance)
 
 
 @receiver(post_auth_failed)
diff --git a/apps/audits/utils.py b/apps/audits/utils.py
index 2c54565d5..44e858098 100644
--- a/apps/audits/utils.py
+++ b/apps/audits/utils.py
@@ -1,12 +1,12 @@
 import copy
-from itertools import chain
 from datetime import datetime
+from itertools import chain
 
 from django.db import models
 
-from common.utils.timezone import as_current_tz
-from common.utils import validate_ip, get_ip_city, get_logger
 from common.db.fields import RelatedManager
+from common.utils import validate_ip, get_ip_city, get_logger
+from common.utils.timezone import as_current_tz
 from .const import DEFAULT_CITY
 
 logger = get_logger(__name__)
@@ -22,7 +22,7 @@ def write_login_log(*args, **kwargs):
     else:
         city = get_ip_city(ip) or DEFAULT_CITY
     kwargs.update({'ip': ip, 'city': city})
-    UserLoginLog.objects.create(**kwargs)
+    return UserLoginLog.objects.create(**kwargs)
 
 
 def _get_instance_field_value(
diff --git a/apps/rbac/builtin.py b/apps/rbac/builtin.py
index 8498f7776..71cd01ff2 100644
--- a/apps/rbac/builtin.py
+++ b/apps/rbac/builtin.py
@@ -36,6 +36,7 @@ system_user_perms += (user_perms + _view_all_joined_org_perms)
 
 _auditor_perms = (
     ('rbac', 'menupermission', 'view', 'audit'),
+    ('users', 'usersession', '*', '*'),
     ('audits', '*', '*', '*'),
     ('audits', 'joblog', '*', '*'),
     ('terminal', 'commandstorage', 'view', 'commandstorage'),
diff --git a/apps/rbac/const.py b/apps/rbac/const.py
index 5825231b5..a63e013e1 100644
--- a/apps/rbac/const.py
+++ b/apps/rbac/const.py
@@ -28,6 +28,7 @@ exclude_permissions = (
     ('authentication', 'superconnectiontoken', 'change,delete', 'superconnectiontoken'),
     ('authentication', 'temptoken', 'delete', 'temptoken'),
     ('users', 'userpasswordhistory', '*', '*'),
+    ('users', 'usersession', 'add,delete,change', 'usersession'),
     ('assets', 'adminuser', '*', '*'),
     ('assets', 'assetgroup', '*', '*'),
     ('assets', 'cluster', '*', '*'),
diff --git a/apps/rbac/tree.py b/apps/rbac/tree.py
index 5bbedbbf0..998f3540f 100644
--- a/apps/rbac/tree.py
+++ b/apps/rbac/tree.py
@@ -98,6 +98,7 @@ special_pid_mapper = {
     'terminal.endpoint': 'terminal_node',
     'terminal.endpointrule': 'terminal_node',
     'audits.ftplog': 'terminal',
+    'users.usersession': 'terminal',
     'perms.view_myassets': 'my_assets',
     'ops.celerytask': 'task_center',
     'ops.view_celerytaskexecution': 'task_center',
diff --git a/apps/terminal/api/session/session.py b/apps/terminal/api/session/session.py
index e4f76435d..756fb5a91 100644
--- a/apps/terminal/api/session/session.py
+++ b/apps/terminal/api/session/session.py
@@ -38,7 +38,7 @@ from users.models import User
 
 __all__ = [
     'SessionViewSet', 'SessionReplayViewSet',
-    'SessionJoinValidateAPI', 'MySessionAPIView',
+    'SessionJoinValidateAPI', 'MySessionAPIView'
 ]
 
 logger = get_logger(__name__)
diff --git a/apps/terminal/models/session/__init__.py b/apps/terminal/models/session/__init__.py
index 073c9d078..255d3c54e 100644
--- a/apps/terminal/models/session/__init__.py
+++ b/apps/terminal/models/session/__init__.py
@@ -1,4 +1,4 @@
 from .command import *
-from .session import *
 from .replay import *
+from .session import *
 from .sharing import *
diff --git a/apps/terminal/signal_handlers/session.py b/apps/terminal/signal_handlers/session.py
index 7aa35307d..c67fcab77 100644
--- a/apps/terminal/signal_handlers/session.py
+++ b/apps/terminal/signal_handlers/session.py
@@ -16,3 +16,4 @@ def on_session_finished(sender, instance: Session, created, **kwargs):
         return
     # 清理一次可能因 task 未执行的缓存数据
     Session.unlock_session(instance.id)
+
diff --git a/apps/users/api/__init__.py b/apps/users/api/__init__.py
index 8553b4bc6..464c857a7 100644
--- a/apps/users/api/__init__.py
+++ b/apps/users/api/__init__.py
@@ -6,4 +6,5 @@ from .preference import *
 from .profile import *
 from .relation import *
 from .service import *
+from .session import *
 from .user import *
diff --git a/apps/users/api/session.py b/apps/users/api/session.py
new file mode 100644
index 000000000..6fb90ea01
--- /dev/null
+++ b/apps/users/api/session.py
@@ -0,0 +1,51 @@
+from importlib import import_module
+
+from django.conf import settings
+from django.utils import timezone
+from rest_framework import status
+from rest_framework import viewsets
+from rest_framework.decorators import action
+from rest_framework.response import Response
+
+from common.api import CommonApiMixin
+from orgs.utils import current_org
+from users import serializers
+from ..models import UserSession
+
+
+class UserSessionViewSet(CommonApiMixin, viewsets.ModelViewSet):
+    http_method_names = ('get', 'post', 'head', 'options', 'trace')
+    serializer_class = serializers.UserSessionSerializer
+    filterset_fields = ['id', 'ip', 'city', 'type']
+    search_fields = ['id', 'ip', 'city']
+
+    rbac_perms = {
+        'offline': ['users.offline_usersession']
+    }
+
+    @property
+    def org_user_ids(self):
+        user_ids = current_org.get_members().values_list('id', flat=True)
+        return user_ids
+
+    def get_queryset(self):
+        queryset = UserSession.objects.filter(date_expired__gt=timezone.now())
+        if current_org.is_root():
+            return queryset
+        user_ids = self.org_user_ids
+        queryset = queryset.filter(user_id__in=user_ids)
+        return queryset
+
+    @action(['POST'], detail=False, url_path='offline')
+    def offline(self, request, *args, **kwargs):
+        ids = request.data.get('ids', [])
+        queryset = self.get_queryset().exclude(key=request.session.session_key).filter(id__in=ids)
+        if not queryset.exists():
+            return Response(status=status.HTTP_200_OK)
+
+        keys = queryset.values_list('key', flat=True)
+        session_store_cls = import_module(settings.SESSION_ENGINE).SessionStore
+        for key in keys:
+            session_store_cls(key).delete()
+        queryset.delete()
+        return Response(status=status.HTTP_200_OK)
diff --git a/apps/users/migrations/0044_usersession.py b/apps/users/migrations/0044_usersession.py
new file mode 100644
index 000000000..325cd4f2c
--- /dev/null
+++ b/apps/users/migrations/0044_usersession.py
@@ -0,0 +1,36 @@
+# Generated by Django 4.1.10 on 2023-09-14 07:23
+
+from django.conf import settings
+from django.db import migrations, models
+import django.db.models.deletion
+import uuid
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ('users', '0043_remove_user_secret_key_preference'),
+    ]
+
+    operations = [
+        migrations.CreateModel(
+            name='UserSession',
+            fields=[
+                ('id', models.UUIDField(default=uuid.uuid4, primary_key=True, serialize=False)),
+                ('ip', models.GenericIPAddressField(verbose_name='Login IP')),
+                ('key', models.CharField(max_length=128, verbose_name='Session key')),
+                ('city', models.CharField(blank=True, max_length=254, null=True, verbose_name='Login city')),
+                ('user_agent', models.CharField(blank=True, max_length=254, null=True, verbose_name='User agent')),
+                ('type', models.CharField(choices=[('W', 'Web'), ('T', 'Terminal'), ('U', 'Unknown')], max_length=2, verbose_name='Login type')),
+                ('backend', models.CharField(default='', max_length=32, verbose_name='Authentication backend')),
+                ('date_created', models.DateTimeField(blank=True, null=True, verbose_name='Date created')),
+                ('date_expired', models.DateTimeField(blank=True, db_index=True, null=True, verbose_name='Date expired')),
+                ('user', models.ForeignKey(on_delete=django.db.models.deletion.CASCADE, related_name='sessions', to=settings.AUTH_USER_MODEL, verbose_name='User')),
+            ],
+            options={
+                'verbose_name': 'User session',
+                'ordering': ['-date_created'],
+                'permissions': [('offline_usersession', 'Offline ussr session')],
+            },
+        ),
+    ]
diff --git a/apps/users/models/__init__.py b/apps/users/models/__init__.py
index 34e081c5a..7f6c9dda6 100644
--- a/apps/users/models/__init__.py
+++ b/apps/users/models/__init__.py
@@ -4,5 +4,6 @@
 
 from .group import *
 from .preference import *
+from .session import *
 from .user import *
 from .utils import *
diff --git a/apps/users/models/session.py b/apps/users/models/session.py
new file mode 100644
index 000000000..eab35ea8e
--- /dev/null
+++ b/apps/users/models/session.py
@@ -0,0 +1,40 @@
+import uuid
+
+from django.db import models
+from django.utils import timezone
+from django.utils.translation import gettext, gettext_lazy as _
+
+from audits.const import LoginTypeChoices
+
+
+class UserSession(models.Model):
+    id = models.UUIDField(default=uuid.uuid4, primary_key=True)
+    ip = models.GenericIPAddressField(verbose_name=_("Login IP"))
+    key = models.CharField(max_length=128, verbose_name=_("Session key"))
+    city = models.CharField(max_length=254, blank=True, null=True, verbose_name=_("Login city"))
+    user_agent = models.CharField(max_length=254, blank=True, null=True, verbose_name=_("User agent"))
+    type = models.CharField(choices=LoginTypeChoices.choices, max_length=2, verbose_name=_("Login type"))
+    backend = models.CharField(max_length=32, default="", verbose_name=_("Authentication backend"))
+    date_created = models.DateTimeField(null=True, blank=True, verbose_name=_('Date created'))
+    date_expired = models.DateTimeField(null=True, blank=True, verbose_name=_("Date expired"), db_index=True)
+    user = models.ForeignKey(
+        'users.User', verbose_name=_('User'), related_name='sessions', on_delete=models.CASCADE
+    )
+
+    def __str__(self):
+        return '%s(%s)' % (self.user, self.ip)
+
+    @property
+    def backend_display(self):
+        return gettext(self.backend)
+
+    @classmethod
+    def clear_expired_sessions(cls):
+        cls.objects.filter(date_expired__lt=timezone.now()).delete()
+
+    class Meta:
+        ordering = ['-date_created']
+        verbose_name = _('User session')
+        permissions = [
+            ('offline_usersession', _('Offline ussr session')),
+        ]
diff --git a/apps/users/serializers/__init__.py b/apps/users/serializers/__init__.py
index 58191b569..82d92d907 100644
--- a/apps/users/serializers/__init__.py
+++ b/apps/users/serializers/__init__.py
@@ -4,4 +4,5 @@ from .group import *
 from .preference import *
 from .profile import *
 from .realtion import *
+from .session import *
 from .user import *
diff --git a/apps/users/serializers/session.py b/apps/users/serializers/session.py
new file mode 100644
index 000000000..8ff3bbec5
--- /dev/null
+++ b/apps/users/serializers/session.py
@@ -0,0 +1,32 @@
+from django.utils.translation import gettext_lazy as _
+from rest_framework import serializers
+
+from audits.const import LoginTypeChoices
+from common.serializers.fields import LabeledChoiceField, ObjectRelatedField
+from users.models import User
+from ..models import UserSession
+
+
+class UserSessionSerializer(serializers.ModelSerializer):
+    type = LabeledChoiceField(choices=LoginTypeChoices.choices, label=_("Type"))
+    user = ObjectRelatedField(required=False, queryset=User.objects, label=_('User'))
+    is_current_user_session = serializers.SerializerMethodField()
+
+    class Meta:
+        model = UserSession
+        fields_mini = ['id']
+        fields_small = fields_mini + [
+            'type', 'ip', 'city', 'user_agent', 'user', 'is_current_user_session',
+            'backend', 'backend_display', 'date_created', 'date_expired'
+        ]
+        fields = fields_small
+        extra_kwargs = {
+            "backend_display": {"label": _("Authentication backend")},
+        }
+
+    def get_is_current_user_session(self, obj):
+        request = self.context.get('request')
+        if not request:
+            return False
+        return request.session.session_key == obj.key
+
diff --git a/apps/users/signal_handlers.py b/apps/users/signal_handlers.py
index 7db379d4a..f84340128 100644
--- a/apps/users/signal_handlers.py
+++ b/apps/users/signal_handlers.py
@@ -1,6 +1,8 @@
 # -*- coding: utf-8 -*-
 #
+from celery import shared_task
 from django.conf import settings
+from django.contrib.auth.signals import user_logged_out
 from django.db.models.signals import post_save
 from django.dispatch import receiver
 from django.utils.translation import gettext_lazy as _
@@ -10,10 +12,12 @@ from django_cas_ng.signals import cas_user_authenticated
 from authentication.backends.oauth2.signals import oauth2_create_or_update_user
 from authentication.backends.oidc.signals import openid_create_or_update_user
 from authentication.backends.saml2.signals import saml2_create_or_update_user
+from common.const.crontab import CRONTAB_AT_PM_TWO
 from common.decorators import on_transaction_commit
 from common.utils import get_logger
 from jumpserver.utils import get_current_request
-from .models import User, UserPasswordHistory
+from ops.celery.decorator import register_as_period_task
+from .models import User, UserPasswordHistory, UserSession
 from .signals import post_user_create
 
 logger = get_logger(__file__)
@@ -156,3 +160,15 @@ def on_openid_create_or_update_user(sender, request, user, created, name, userna
         user.username = username
         user.email = email
         user.save()
+
+
+@shared_task(verbose_name=_('Clean audits session task log'))
+@register_as_period_task(crontab=CRONTAB_AT_PM_TWO)
+def clean_audits_log_period():
+    UserSession.clear_expired_sessions()
+
+
+@receiver(user_logged_out)
+def user_logged_out_callback(sender, request, user, **kwargs):
+    session_key = request.session.session_key
+    UserSession.objects.filter(key=session_key).delete()
diff --git a/apps/users/urls/api_urls.py b/apps/users/urls/api_urls.py
index fa44e67e4..9f8247a3f 100644
--- a/apps/users/urls/api_urls.py
+++ b/apps/users/urls/api_urls.py
@@ -17,6 +17,7 @@ router.register(r'groups', api.UserGroupViewSet, 'user-group')
 router.register(r'users-groups-relations', api.UserUserGroupRelationViewSet, 'users-groups-relation')
 router.register(r'service-account-registrations', api.ServiceAccountRegistrationViewSet, 'service-account-registration')
 router.register(r'connection-token', auth_api.ConnectionTokenViewSet, 'connection-token')
+router.register(r'user-sessions', api.UserSessionViewSet, 'user-session')
 
 
 urlpatterns = [