# -*- coding: utf-8 -*- # from collections import OrderedDict import logging import os import uuid from django.core.cache import cache from django.shortcuts import get_object_or_404, redirect from django.utils import timezone from django.core.files.storage import default_storage from django.http.response import HttpResponseRedirectBase from django.http import HttpResponseNotFound from django.conf import settings import jms_storage from rest_framework import viewsets, serializers from rest_framework.views import APIView, Response from rest_framework.permissions import AllowAny from rest_framework_bulk import BulkModelViewSet from common.utils import get_object_or_none from .models import Terminal, Status, Session, Task from .serializers import TerminalSerializer, StatusSerializer, \ SessionSerializer, TaskSerializer, ReplaySerializer from .hands import IsSuperUserOrAppUser, IsAppUser, \ IsSuperUserOrAppUserOrUserReadonly from .backends import get_command_storage, get_multi_command_storage, \ SessionCommandSerializer logger = logging.getLogger(__file__) class TerminalViewSet(viewsets.ModelViewSet): queryset = Terminal.objects.filter(is_deleted=False) serializer_class = TerminalSerializer permission_classes = (IsSuperUserOrAppUserOrUserReadonly,) def create(self, request, *args, **kwargs): name = request.data.get('name') remote_ip = request.META.get('REMOTE_ADDR') x_real_ip = request.META.get('X-Real-IP') remote_addr = x_real_ip or remote_ip terminal = get_object_or_none(Terminal, name=name, is_deleted=False) if terminal: msg = 'Terminal name %s already used' % name return Response({'msg': msg}, status=409) serializer = self.serializer_class(data={ 'name': name, 'remote_addr': remote_addr }) if serializer.is_valid(): terminal = serializer.save() # App should use id, token get access key, if accepted token = uuid.uuid4().hex cache.set(token, str(terminal.id), 3600) data = {"id": str(terminal.id), "token": token, "msg": "Need accept"} return Response(data, status=201) else: data = serializer.errors logger.error("Register terminal error: {}".format(data)) return Response(data, status=400) def get_permissions(self): if self.action == "create": self.permission_classes = (AllowAny,) return super().get_permissions() class TerminalTokenApi(APIView): permission_classes = (AllowAny,) queryset = Terminal.objects.filter(is_deleted=False) def get(self, request, *args, **kwargs): try: terminal = self.queryset.get(id=kwargs.get('terminal')) except Terminal.DoesNotExist: terminal = None token = request.query_params.get("token") if terminal is None: return Response('May be reject by administrator', status=401) if token is None or cache.get(token, "") != str(terminal.id): return Response('Token is not valid', status=401) if not terminal.is_accepted: return Response("Terminal was not accepted yet", status=400) if not terminal.user or not terminal.user.access_key.all(): return Response("No access key generate", status=401) access_key = terminal.user.access_key.first() data = OrderedDict() data['access_key'] = {'id': access_key.id, 'secret': access_key.secret} return Response(data, status=200) class StatusViewSet(viewsets.ModelViewSet): queryset = Status.objects.all() serializer_class = StatusSerializer permission_classes = (IsSuperUserOrAppUser,) session_serializer_class = SessionSerializer task_serializer_class = TaskSerializer def create(self, request, *args, **kwargs): from_gua = self.request.query_params.get("from_guacamole", None) if not from_gua: self.handle_sessions() super().create(request, *args, **kwargs) tasks = self.request.user.terminal.task_set.filter(is_finished=False) serializer = self.task_serializer_class(tasks, many=True) return Response(serializer.data, status=201) def handle_sessions(self): sessions_active = [] for session_data in self.request.data.get("sessions", []): self.create_or_update_session(session_data) if not session_data["is_finished"]: sessions_active.append(session_data["id"]) sessions_in_db_active = Session.objects.filter( is_finished=False, terminal=self.request.user.terminal.id ) for session in sessions_in_db_active: if str(session.id) not in sessions_active: session.is_finished = True session.date_end = timezone.now() session.save() def create_or_update_session(self, session_data): session_data["terminal"] = self.request.user.terminal.id _id = session_data["id"] session = get_object_or_none(Session, id=_id) if session: serializer = SessionSerializer( data=session_data, instance=session ) else: serializer = SessionSerializer(data=session_data) if serializer.is_valid(): session = serializer.save() return session else: msg = "session data is not valid {}: {}".format( serializer.errors, str(serializer.data) ) logger.error(msg) return None def get_queryset(self): terminal_id = self.kwargs.get("terminal", None) if terminal_id: terminal = get_object_or_404(Terminal, id=terminal_id) self.queryset = terminal.status_set.all() return self.queryset def perform_create(self, serializer): serializer.validated_data["terminal"] = self.request.user.terminal return super().perform_create(serializer) def get_permissions(self): if self.action == "create": self.permission_classes = (IsAppUser,) return super().get_permissions() class SessionViewSet(viewsets.ModelViewSet): queryset = Session.objects.all() serializer_class = SessionSerializer permission_classes = (IsSuperUserOrAppUser,) def get_queryset(self): terminal_id = self.kwargs.get("terminal", None) if terminal_id: terminal = get_object_or_404(Terminal, id=terminal_id) self.queryset = terminal.session_set.all() return self.queryset def perform_create(self, serializer): if self.request.user.terminal: serializer.validated_data["terminal"] = self.request.user.terminal return super().perform_create(serializer) class TaskViewSet(BulkModelViewSet): queryset = Task.objects.all() serializer_class = TaskSerializer permission_classes = (IsSuperUserOrAppUser,) class KillSessionAPI(APIView): permission_classes = (IsSuperUserOrAppUser,) model = Task def post(self, request, *args, **kwargs): validated_session = [] for session_id in request.data: session = get_object_or_none(Session, id=session_id) if session and not session.is_finished: validated_session.append(session_id) self.model.objects.create( name="kill_session", args=session.id, terminal=session.terminal, ) return Response({"ok": validated_session}) class CommandViewSet(viewsets.ViewSet): """接受app发送来的command log, 格式如下 { "user": "admin", "asset": "localhost", "system_user": "web", "session": "xxxxxx", "input": "whoami", "output": "d2hvbWFp", # base64.b64encode(s) "timestamp": 1485238673.0 } """ command_store = get_command_storage() multi_command_storage = get_multi_command_storage() serializer_class = SessionCommandSerializer permission_classes = (IsSuperUserOrAppUser,) def get_queryset(self): self.command_store.filter(**dict(self.request.query_params)) def create(self, request, *args, **kwargs): serializer = self.serializer_class(data=request.data, many=True) if serializer.is_valid(): ok = self.command_store.bulk_save(serializer.validated_data) if ok: return Response("ok", status=201) else: return Response("Save error", status=500) else: msg = "Not valid: {}".format(serializer.errors) logger.error(msg) return Response({"msg": msg}, status=401) def list(self, request, *args, **kwargs): queryset = self.multi_command_storage.filter() serializer = self.serializer_class(queryset, many=True) return Response(serializer.data) class SessionReplayViewSet(viewsets.ViewSet): serializer_class = ReplaySerializer permission_classes = (IsSuperUserOrAppUser,) session = None upload_to = 'replay' # 仅添加到本地存储中 def get_session_path(self, version=2): """ 获取session日志的文件路径 :param version: 原来后缀是 .gz，为了统一新版本改为 .replay.gz :return: """ suffix = '.replay.gz' if version == 1: suffix = '.gz' date = self.session.date_start.strftime('%Y-%m-%d') return os.path.join(date, str(self.session.id) + suffix) def get_local_path(self, version=2): session_path = self.get_session_path(version=version) if version == 2: local_path = os.path.join(self.upload_to, session_path) else: local_path = session_path return local_path def save_to_storage(self, f): local_path = self.get_local_path() try: name = default_storage.save(local_path, f) return name, None except OSError as e: return None, e def create(self, request, *args, **kwargs): session_id = kwargs.get('pk') self.session = get_object_or_404(Session, id=session_id) serializer = self.serializer_class(data=request.data) if serializer.is_valid(): file = serializer.validated_data['file'] name, err = self.save_to_storage(file) if not name: msg = "Failed save replay `{}`: {}".format(session_id, err) logger.error(msg) return Response({'msg': str(err)}, status=400) url = default_storage.url(name) return Response({'url': url}, status=201) else: msg = 'Upload data invalid: {}'.format(serializer.errors) logger.error(msg) return Response({'msg': serializer.errors}, status=401) def retrieve(self, request, *args, **kwargs): session_id = kwargs.get('pk') self.session = get_object_or_404(Session, id=session_id) # 新版本和老版本的文件后缀不同 session_path = self.get_session_path() # 存在外部存储上的路径 local_path = self.get_local_path() local_path_v1 = self.get_local_path(version=1) # 去default storage中查找 for _local_path in (local_path, local_path_v1, session_path): if default_storage.exists(_local_path): url = default_storage.url(_local_path) return redirect(url) # 去定义的外部storage查找 configs = settings.TERMINAL_REPLAY_STORAGE configs = {k: v for k, v in configs.items() if v['TYPE'] != 'server'} if not configs: return HttpResponseNotFound() target_path = os.path.join(default_storage.base_location, local_path) # 保存到storage的路径 target_dir = os.path.dirname(target_path) if not os.path.isdir(target_dir): os.makedirs(target_dir, exist_ok=True) storage = jms_storage.get_multi_object_storage(configs) ok, err = storage.download(session_path, target_path) if not ok: logger.error("Failed download replay file: {}".format(err)) return HttpResponseNotFound() return redirect(default_storage.url(local_path)) class SessionReplayV2ViewSet(SessionReplayViewSet): serializer_class = ReplaySerializer permission_classes = (IsSuperUserOrAppUser,) session = None def retrieve(self, request, *args, **kwargs): response = super().retrieve(request, *args, **kwargs) data = { 'type': 'guacamole' if self.session.protocol == 'rdp' else 'json', 'src': '', } if isinstance(response, HttpResponseRedirectBase): data['src'] = response.url return Response(data) return HttpResponseNotFound() class TerminalConfig(APIView): permission_classes = (IsAppUser,) def get(self, request): user = request.user terminal = user.terminal configs = terminal.config return Response(configs, status=200)