aes_field.py 3.48 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 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67
from django.conf import settings

EOD = '`%EofD%`'  # This should be something that will not occur in strings


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)])


def encrypt(key, s):
    """ AES Encrypt a secret `s` with the key `key` """
    obj = AES.new(key)
    datalength = len(s) + len(EOD)
    if datalength < 16:
        saltlength = 16 - datalength
    else:
        saltlength = 16 - datalength % 16
    ss = ''.join([s, EOD, genstring(saltlength)])
    return obj.encrypt(ss)


def decrypt(key, s):
    """ AES Decrypt a secret `s` with the key `key` """
    obj = AES.new(key)
    ss = obj.decrypt(s)
    return ss.split(bytes(EOD, 'utf-8'))[0]


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 84 85 86
        try:
            return decrypt(settings.AES_KEY,
                           binascii.a2b_base64(value)).decode('utf-8')
        except Exception as e:
87
            raise ValueError(value)
88 89 90 91

    def from_db_value(self, value, *args, **kwargs):
        if value is None:
            return value
92 93
        try:
            return decrypt(settings.AES_KEY,
94
                           binascii.a2b_base64(value)).decode('utf-8')
95
        except Exception as e:
96
            raise ValueError(value)
97 98 99 100 101 102 103

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

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