jumpserver/apps/accounts/models/mixins/vault.py

from django.db import models
from django.db.models.signals import post_save
from django.utils.translation import gettext_lazy as _

from common.db import fields

__all__ = ['VaultQuerySetMixin', 'VaultManagerMixin', 'VaultModelMixin']


class VaultQuerySetMixin(models.QuerySet):

    def update(self, **kwargs):
        """
           1. 替换 secret 为 _secret
           2. 触发 post_save 信号
        """
        if 'secret' in kwargs:
            kwargs.update({
                '_secret': kwargs.pop('secret')
            })
        rows = super().update(**kwargs)

        # 为了获取更新后的对象所以单独查询一次
        ids = self.values_list('id', flat=True)
        objs = self.model.objects.filter(id__in=ids)
        for obj in objs:
            post_save.send(obj.__class__, instance=obj, created=False)
        return rows


class VaultManagerMixin(models.Manager):
    """ 触发 bulk_create 和 bulk_update 操作下的 post_save 信号 """

    def bulk_create(self, objs, batch_size=None, ignore_conflicts=False):
        objs = super().bulk_create(objs, batch_size=batch_size, ignore_conflicts=ignore_conflicts)
        for obj in objs:
            post_save.send(obj.__class__, instance=obj, created=True)
        return objs

    def bulk_update(self, objs, batch_size=None, ignore_conflicts=False):
        objs = super().bulk_update(objs, batch_size=batch_size, ignore_conflicts=ignore_conflicts)
        for obj in objs:
            post_save.send(obj.__class__, instance=obj, created=False)
        return objs


class VaultModelMixin(models.Model):
    _secret = fields.EncryptTextField(blank=True, null=True, verbose_name=_('Secret'))
    is_sync_metadata = True

    class Meta:
        abstract = True

    # 缓存 secret 值, lazy-property 不能用
    __secret = None

    @property
    def secret(self):
        if self.__secret:
            return self.__secret
        from accounts.backends import vault_client
        secret = vault_client.get(self)
        if not secret and not self.secret_has_save_to_vault:
            # vault_client 获取不到, 并且 secret 没有保存到 vault, 就从 self._secret 获取
            secret = self._secret
        self.__secret = secret
        return self.__secret

    @secret.setter
    def secret(self, value):
        """
        保存的时候通过 post_save 信号监听进行处理,
        先保存到 db, 再保存到 vault 同时删除本地 db _secret 值
        """
        self._secret = value
        self.__secret = value

    _secret_save_to_vault_mark = '# Secret-has-been-saved-to-vault #'

    def mark_secret_save_to_vault(self):
        self._secret = self._secret_save_to_vault_mark
        self.save()

    @property
    def secret_has_save_to_vault(self):
        return self._secret == self._secret_save_to_vault_mark

    def save(self, *args, **kwargs):
        """ 通过 post_save signal 处理 _secret 数据 """
        update_fields = kwargs.get('update_fields')
        if update_fields and 'secret' in update_fields:
            update_fields.remove('secret')
            update_fields.append('_secret')
        return super().save(*args, **kwargs)