khaganat-web/pwdb/models.py

from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
from cryptography.hazmat.primitives.kdf.hkdf import HKDF
from cryptography.hazmat.primitives import hashes, padding
from cryptography.hazmat.backends import default_backend
from django.utils.translation import ugettext_lazy as _
from django.utils.encoding import force_bytes
from django.conf import settings
from django.db import models
from neluser.models import NelUser
import secrets
import uuid

KEY_LENGTH = 32
IV_LENGTH = 16
BLOCK_SIZE = 128
ENCODING = "UTF-8"


class SharedPassword(models.Model):
    uuid = models.UUIDField(primary_key=True, default=uuid.uuid4, editable=False)
    name = models.CharField(max_length=512)
    url = models.CharField(max_length=512, blank=True)
    description = models.TextField(blank=True)
    iv = models.BinaryField(max_length=IV_LENGTH)
    encrypted_password = models.BinaryField(max_length=2048)

    @staticmethod
    def get_key(pass_uuid, key=None):
        backend = default_backend()
        hkdf = HKDF(
            algorithm=hashes.SHA256(),
            length=KEY_LENGTH,
            salt=pass_uuid.bytes,
            backend=backend,
            info=None,
        )
        key = key or settings.SECRET_KEY
        key = force_bytes(key, encoding=ENCODING)
        return hkdf.derive(key)

    @staticmethod
    def get_cipher(key, iv):
        backend = default_backend()
        return Cipher(algorithms.Camellia(key), modes.CBC(iv), backend=backend)

    @staticmethod
    def padd_password(clear_password):
        clear_password = force_bytes(clear_password, encoding=ENCODING)
        padder = padding.PKCS7(BLOCK_SIZE).padder()
        padded_password = padder.update(clear_password) + padder.finalize()
        return padded_password

    @staticmethod
    def unpadd_password(clear_password):
        unpadder = padding.PKCS7(BLOCK_SIZE).unpadder()
        unpadded_password = unpadder.update(clear_password) + unpadder.finalize()
        return unpadded_password.decode()

    @staticmethod
    def new(name, clear_password):
        password = SharedPassword(
            uuid=uuid.uuid4(),
            name=name,
            url="",
            description="",
            iv=secrets.token_bytes(IV_LENGTH),
            encrypted_password=b"",
        )
        password.set_password(clear_password)
        return password

    def set_password(self, clear_password):
        clear_password = SharedPassword.padd_password(clear_password)
        key = SharedPassword.get_key(self.uuid)
        cipher = SharedPassword.get_cipher(key, self.iv)
        encryptor = cipher.encryptor()
        self.encrypted_password = (
            encryptor.update(clear_password) + encryptor.finalize()
        )

    def decrypt_password(self, key=None):
        key = SharedPassword.get_key(self.uuid, key=key)
        cipher = SharedPassword.get_cipher(key, self.iv)
        decryptor = cipher.decryptor()
        clear_password = (
            decryptor.update(self.encrypted_password) + decryptor.finalize()
        )
        clear_password = SharedPassword.unpadd_password(clear_password)
        return clear_password

    def users(self):
        lst = (
            SharedPasswordAccess.objects.select_related("user")
            .filter(password=self)
            .order_by("user__email")
        )
        return [e.user for e in lst]

    def __str__(self):
        return self.name

    class Meta:
        verbose_name = _("shared_password")
        verbose_name_plural = _("shared_passwords")


class SharedPasswordAccess(models.Model):
    user = models.ForeignKey(
        NelUser, on_delete=models.CASCADE, limit_choices_to={"is_staff": True}
    )
    password = models.ForeignKey(SharedPassword, on_delete=models.CASCADE)

    def __str__(self):
        return "{}: {}".format(self.password, self.user)

    class Meta:
        verbose_name = _("shared_password_access")
        verbose_name_plural = _("shared_passwords_access")