aes_field.py 3.76 KB
Newer Older
1 2 3 4 5 6 7 8 9
# coding:utf-8
# Re2o est un logiciel d'administration développé initiallement au rezometz. Il
# se veut agnostique au réseau considéré, de manière à être installable en
# quelques clics.
#
# Copyright © 2017  Gabriel Détraz
# Copyright © 2017  Goulven Kermarec
# Copyright © 2017  Augustin Lemesle
# Copyright © 2018  Maël Kervella
10
# Copyright © 2018  Hugo Levy-Falk
11 12 13 14 15 16 17 18 19 20 21 22 23 24 25
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation; either version 2 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License along
# with this program; if not, write to the Free Software Foundation, Inc.,
# 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.

26
"""
27 28 29 30 31 32 33 34 35 36
Module defining a AESEncryptedField object that can be used in forms
to handle the use of properly encrypting and decrypting AES keys
"""

import string
import binascii
from random import choice
from Crypto.Cipher import AES

from django.db import models
37
from django import forms
38 39
from django.conf import settings

40 41
EOD_asbyte = b'`%EofD%`'  # This should be something that will not occur in strings
EOD = EOD_asbyte.decode('utf-8')
42 43 44 45 46 47 48

def genstring(length=16, chars=string.printable):
    """ Generate a random string of length `length` and composed of
    the characters in `chars` """
    return ''.join([choice(chars) for i in range(length)])


49 50
def encrypt(key, secret):
    """ AES Encrypt a secret with the key `key` """
51
    obj = AES.new(key)
52
    datalength = len(secret) + len(EOD)
53 54 55 56
    if datalength < 16:
        saltlength = 16 - datalength
    else:
        saltlength = 16 - datalength % 16
57 58
    encrypted_secret = ''.join([secret, EOD, genstring(saltlength)])
    return obj.encrypt(encrypted_secret)
59 60


61 62
def decrypt(key, secret):
    """ AES Decrypt a secret with the key `key` """
63
    obj = AES.new(key)
64 65
    uncrypted_secret = obj.decrypt(secret)
    return uncrypted_secret.split(EOD_asbyte)[0]
66 67


68 69 70 71
class AESEncryptedFormField(forms.CharField):
    widget = forms.PasswordInput(render_value=True)


72 73 74
class AESEncryptedField(models.CharField):
    """ A Field that can be used in forms for adding the support
    of AES ecnrypted fields """
75

76
    def save_form_data(self, instance, data):
77 78
        setattr(instance, self.name, binascii.b2a_base64(
            encrypt(settings.AES_KEY, data)).decode('utf-8'))
79 80 81 82

    def to_python(self, value):
        if value is None:
            return None
83
        try:
84
            return decrypt(settings.AES_KEY, binascii.a2b_base64(value)).decode('utf-8')
85 86 87 88 89
        except UnicodeDecodeError as e:
            raise ValueError(
                "Could not decode your field %s, your settings.AES_KEY "
                "is probably wrong." % self.name
            )
90 91 92 93

    def from_db_value(self, value, *args, **kwargs):
        if value is None:
            return value
94
        try:
95
            return decrypt(settings.AES_KEY, binascii.a2b_base64(value)).decode('utf-8')
96 97 98 99 100
        except UnicodeDecodeError as e:
            raise ValueError(
                "Could not decode your field %s, your settings.AES_KEY "
                "is probably wrong." % self.name
            )
101 102 103 104

    def get_prep_value(self, value):
        if value is None:
            return value
105
        return binascii.b2a_base64(encrypt(settings.AES_KEY, value)).decode('utf-8')
106 107 108 109 110

    def formfield(self, **kwargs):
        defaults = {'form_class': AESEncryptedFormField}
        defaults.update(kwargs)
        return super().formfield(**defaults)