jumpserver/apps/common/drf/renders/base.py

import abc
import io
import re
from datetime import datetime

import pyzipper
from django.utils.translation import gettext_lazy as _
from rest_framework import serializers
from rest_framework.renderers import BaseRenderer
from rest_framework.utils import encoders, json

from common.serializers.fields import ObjectRelatedField, LabeledChoiceField
from common.utils import get_logger

logger = get_logger(__file__)


class BaseFileRenderer(BaseRenderer):
    # 渲染模版标识, 导入、导出、更新模版: ['import', 'update', 'export']
    template = 'export'
    serializer = None

    @staticmethod
    def _check_validation_data(data):
        detail_key = "detail"
        if detail_key in data:
            return False
        return True

    @staticmethod
    def _json_format_response(response_data):
        return json.dumps(response_data)

    def set_response_disposition(self, response):
        serializer = self.serializer
        if response and hasattr(serializer, 'Meta') and hasattr(serializer.Meta, "model"):
            filename_prefix = serializer.Meta.model.__name__.lower()
        else:
            filename_prefix = 'download'
        now = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
        filename = "{}_{}.{}".format(filename_prefix, now, self.format)
        disposition = 'attachment; filename="{}"'.format(filename)
        response['Content-Disposition'] = disposition

    def get_rendered_fields(self):
        fields = self.serializer.fields
        if self.template == 'import':
            fields = [v for k, v in fields.items() if not v.read_only and k != "org_id" and k != 'id']
        elif self.template == 'update':
            fields = [v for k, v in fields.items() if not v.read_only and k != "org_id"]
        else:
            fields = [v for k, v in fields.items() if not v.write_only and k != "org_id"]

        meta = getattr(self.serializer, 'Meta', None)
        if meta:
            fields_unexport = getattr(meta, 'fields_unexport', [])
            fields = [v for v in fields if v.field_name not in fields_unexport]
        return fields

    @staticmethod
    def get_column_titles(render_fields):
        titles = []
        for field in render_fields:
            name = field.label
            if field.required:
                name = '*' + name
            titles.append(name)
        return titles

    def process_data(self, data):
        results = data['results'] if 'results' in data else data

        if isinstance(results, dict):
            results = [results]

        if self.template == 'import':
            results = [results[0]] if results else results
        else:
            # 限制数据数量
            results = results[:10000]
        # 会将一些 UUID 字段转化为 string
        results = json.loads(json.dumps(results, cls=encoders.JSONEncoder))
        return results

    @staticmethod
    def to_id_name(value):
        if value is None:
            return '-'
        pk = str(value.get('id', '') or value.get('pk', ''))
        name = value.get('name') or value.get('display_name', '')
        return '{}({})'.format(name, pk)

    @staticmethod
    def to_choice_name(value):
        if value is None:
            return '-'
        value = value.get('value', '')
        return value

    def render_value(self, field, value):
        if value is None:
            value = '-'
        elif hasattr(field, 'to_file_representation'):
            value = field.to_file_representation(value)
        elif isinstance(value, bool):
            value = 'Yes' if value else 'No'
        elif isinstance(field, LabeledChoiceField):
            value = value.get('value', '')
        elif isinstance(field, ObjectRelatedField):
            if field.many:
                value = [self.to_id_name(v) for v in value]
            else:
                value = self.to_id_name(value)
        elif isinstance(field, serializers.ListSerializer):
            value = [self.render_value(field.child, v) for v in value]
        elif isinstance(field, serializers.Serializer) and value.get('id'):
            value = self.to_id_name(value)
        elif isinstance(field, serializers.ManyRelatedField):
            value = [self.render_value(field.child_relation, v) for v in value]
        elif isinstance(field, serializers.ListField):
            value = [self.render_value(field.child, v) for v in value]

        if not isinstance(value, str):
            value = json.dumps(value, cls=encoders.JSONEncoder, ensure_ascii=False)
        return str(value)

    def generate_rows(self, data, render_fields):
        for item in data:
            row = []
            for field in render_fields:
                value = item.get(field.field_name)
                value = self.render_value(field, value)
                row.append(value)
            yield row

    @abc.abstractmethod
    def initial_writer(self):
        raise NotImplementedError

    def write_column_titles(self, column_titles):
        self.write_row(column_titles)

    def write_rows(self, rows):
        for row in rows:
            self.write_row(row)

    @abc.abstractmethod
    def write_row(self, row):
        raise NotImplementedError

    @abc.abstractmethod
    def get_rendered_value(self):
        raise NotImplementedError

    def after_render(self):
        pass

    def render(self, data, accepted_media_type=None, renderer_context=None):
        if data is None:
            return bytes()

        if not self._check_validation_data(data):
            return self._json_format_response(data)

        try:
            renderer_context = renderer_context or {}
            request = renderer_context['request']
            response = renderer_context['response']
            view = renderer_context['view']
            self.template = request.query_params.get('template', 'export')
            self.serializer = view.get_serializer()
            self.set_response_disposition(response)
        except Exception as e:
            logger.debug(e, exc_info=True)
            value = 'The resource not support export!'.encode('utf-8')
            return value

        try:
            rendered_fields = self.get_rendered_fields()
            column_titles = self.get_column_titles(rendered_fields)
            data = self.process_data(data)
            rows = self.generate_rows(data, rendered_fields)
            self.initial_writer()
            self.write_column_titles(column_titles)
            self.write_rows(rows)
            self.after_render()
            value = self.get_rendered_value()
            if getattr(view, 'export_as_zip', False) and self.template == 'export':
                value = self.compress_into_zip_file(value, request, response)
        except Exception as e:
            logger.debug(e, exc_info=True)
            value = 'Render error! ({})'.format(self.media_type).encode('utf-8')
            return value
        return value

    def compress_into_zip_file(self, value, request, response):
        filename_pattern = re.compile(r'filename="([^"]+)"')
        content_disposition = response['Content-Disposition']
        match = filename_pattern.search(content_disposition)
        filename = match.group(1)
        response['Content-Disposition'] = content_disposition.replace(self.format, 'zip')

        contents_io = io.BytesIO()
        secret_key = request.user.secret_key
        if not secret_key:
            content = _("{} - The encryption password has not been set - "
                        "please go to personal information -> file encryption password "
                        "to set the encryption password").format(request.user.name)

            response['Content-Disposition'] = content_disposition.replace(self.format, 'txt')
            contents_io.write(content.encode('utf-8'))
            return contents_io.getvalue()

        with pyzipper.AESZipFile(
                contents_io, 'w', compression=pyzipper.ZIP_LZMA, encryption=pyzipper.WZ_AES
        ) as zf:
            zf.setpassword(secret_key.encode('utf8'))
            zf.writestr(filename, value)
        return contents_io.getvalue()