Irfan Harun
Irfan Harun

Reputation: 1059

unable to write test for django model containing custom model field

I have a model for users where in the field for password is a custom field. This model works fine but i'm not able to run tests for the model.

my model

from core_engine.CustomModelFields import *

class Users(models.Model):
    username = models.EmailField(unique=True)
    name = models.CharField(max_length=100)
    password = EncryptedField(max_length=500)

in my core_engine.CustomModelFields.py file

from account_engine.crypto import *

class EncryptedField(CharField):

    def from_db_value(self, value, expression, connection, context):
        value = Password.decrypt(str(value))
        return value.decode("utf-8")

    def to_python(self, value):
        if not value:
            return None
        return value

        #Only need to decrypt if password already encrypted. 
        try:
            if Password.encrypt(Password.decrypt(value)) == value:
                value = Password.decrypt(str(value))
                return value.decode("utf-8")
        except:
            return value

    def get_db_prep_save(self, value, connection):
        value = Password.encrypt(str(value))
        return value    

and finally in accounts_engine.crypto.py file i have

import base64, hashlib
from django.db import models
from Crypto import Random
from Crypto.Cipher import AES

BS = 16
pad = lambda s: s + (BS - len(s) % BS) * chr(BS - len(s) % BS)
unpad = lambda s : s[0:-s[-1]]


class Password(models.Model):



    def encrypt( raw ):
        mysecretpassword = 'somepassword'
        KEY = hashlib.sha256(mysecretpassword.encode('utf-8')).digest()
        raw = pad(raw)
        iv = Random.new().read( AES.block_size )
        cipher = AES.new( KEY, AES.MODE_CBC, iv )
        return base64.b64encode( iv + cipher.encrypt( raw ) )

    def decrypt( enc ):
        mysecretpassword = 'somepassword'
        KEY = hashlib.sha256(mysecretpassword.encode('utf-8')).digest()  
        enc = base64.b64decode(enc)
        iv = enc[:16]
        cipher = AES.new(KEY, AES.MODE_CBC, iv )
        return unpad(cipher.decrypt( enc[16:] ))

All i want to do is just test my Users model and see that i'm able to create a user during test, which will be required for other tests

so in my test.py file i have

class UsersTestCase(TestCase):

    @classmethod
    def setUp(self):
        print(dt.datetime.now())
        self.user= Users.objects.create(
            username = '[email protected]',
            date_first_registered = dt.datetime.now(),
            password = Password.encrypt('abc')
        )

    def test_get_user(self):

        first_customer = Users.objects.first()
        self.assertEqual(first_customer.username, '[email protected]')

On running the above test, i get an error stating :

TypeError: Object type <class 'str'> cannot be passed to C code

Edit : i understand that the error is due to me passing the password as Password.encrypt('abc').

what changes should i make to my test function in order to create to create a new user

TRACEBACK

Traceback (most recent call last):
  File "D:\project_path\account_engine\tests\tests_models.py", line 15, in setUp
    password = Password.encrypt('abc')
  File "D:\project_path\account_engine\crypto.py", line 21, in encrypt
    return base64.b64encode( iv + cipher.encrypt( raw ) )
  File "d:\project_path\venv\lib\site-packages\Crypto\Cipher\_mode_cbc.py", line 178, in encrypt
    c_uint8_ptr(plaintext),
  File "d:\project_path\venv\lib\site-packages\Crypto\Util\_raw_api.py", line 144, in c_uint8_ptr
    raise TypeError("Object type %s cannot be passed to C code" % type(data))
TypeError: Object type <class 'str'> cannot be passed to C code

Upvotes: 1

Views: 65

Answers (1)

Charnel
Charnel

Reputation: 4432

Encrypt method accepts bytes string type and not str in Python 3:

Plaintexts and ciphertexts (input/output) can only be bytes, bytearray or memoryview. In Python 3, you cannot pass strings. In Python 2, you cannot pass Unicode strings.

You need to encode raw first:

def encrypt( raw ):
    mysecretpassword = 'somepassword'
    KEY = hashlib.sha256(mysecretpassword.encode('utf-8')).digest()
    raw = pad(raw)
    iv = Random.new().read( AES.block_size )
    cipher = AES.new( KEY, AES.MODE_CBC, iv )
    return base64.b64encode(iv + cipher.encrypt(raw.encode('utf-8')))

Upvotes: 1

Related Questions