I'm writing a password manager application, mostly with the goal of learning about proper handling of encryption.
Here is my code responsible for encrypting/decrypting passwords and loading to/from file. The PasswordDataBase class is pretty bare as is, but will incorporate more functionality going further (such as filtering records based on name or tags). I just want to make sure this more sensitive aspect of the code is done right before building upon it.
Dependencies: this code depends on the pyaes module, which you can find on PyPI.
passworddb.py:
from dataclasses import dataclass, field
from hashlib import scrypt
import pickle
import secrets
from typing import List
import uuid
import pyaes
class InvalidPasswordException(Exception):
pass
class PasswordDataBase:
SALT_LENGTH = 16
@dataclass
class Record:
name: str = field(default='')
login: str = field(default='')
password: str = field(default='')
uid: uuid.UUID = field(default_factory=uuid.uuid4)
tags: List[str] = field(default_factory=list)
def __init__(self):
self.records = []
def save_to_file(self, path: str, password: str):
"""
Save the current instance to a file.
Parameters
----------
path : str
Path of the file to save.
password : str
The master password used for encryption.
Returns
-------
None.
"""
password = bytes(password, encoding='utf-8')
salt = secrets.token_bytes(self.SALT_LENGTH)
key = scrypt(password,
n=2**14,
r=8,
p=1,
salt=salt,
dklen=32)
cipher = pyaes.AESModeOfOperationCTR(key)
with open(path, 'wb') as file:
file.write(salt)
file.write(cipher.encrypt(pickle.dumps(self)))
@classmethod
def load_from_file(cls, path: str, password: str):
"""
Load an encrypted password database from a file.
Parameters
----------
path : str
Path to the file to load.
password : str
Master password used to decrypt the file.
Raises
------
InvalidPasswordException
The provided password was invalid.
Returns
-------
PasswordDataBase
A PasswordDatabase instance containing the decrypted data.
"""
password = bytes(password, encoding='utf-8')
with open(path, 'rb') as file:
salt = file.read(cls.SALT_LENGTH)
data = file.read()
key = scrypt(password,
n=2**14,
r=8,
p=1,
salt=salt,
dklen=32)
cipher = pyaes.AESModeOfOperationCTR(key)
try:
return pickle.loads(cipher.decrypt(data))
except pickle.UnpicklingError as error:
raise InvalidPasswordException("Invalid password") from error
def generate_password(length: int, charset: str) -> str:
"""
Generate a securely random password.
Parameters
----------
length : int
Generated password length.
charset : str
Allowable characters for password generation.
Returns
-------
str
The generated password.
"""
return ''.join(secrets.choice(charset) for _ in range(length))
passworddb_tests.py:
"""
Unit tests for the password_manager modules
"""
import os
from string import ascii_letters, digits, punctuation
import tempfile
import unittest
from passworddb import (PasswordDataBase,
InvalidPasswordException,
generate_password)
ALL_CHARS = ascii_letters + digits + punctuation
class PasswordDataBaseTests(unittest.TestCase):
def test_roundtrip(self):
"""
Verify that a password database written to a file is correctly
recovered upon loading.
"""
original = PasswordDataBase()
original.records.append(
PasswordDataBase.Record(name='foo',
login='bar',
password=generate_password(64, ALL_CHARS)))
original.records.append(
PasswordDataBase.Record(name='ga',
login='bu',
password=generate_password(64, ALL_CHARS),
tags=['zo', 'meu']))
original.records.append(PasswordDataBase.Record())
password = "Secr3t"
with tempfile.TemporaryDirectory() as temp_dir:
path = os.path.join(temp_dir, 'test')
original.save_to_file(path, password)
recovered = PasswordDataBase.load_from_file(path, password)
for expected, actual in zip(original.records, recovered.records):
self.assertEqual(expected, actual)
def test_roundtrip_empty(self):
"""
Verify that an empty password database written to a file is correctly
recovered upon loading.
"""
original = PasswordDataBase()
password = "Secr3t"
with tempfile.TemporaryDirectory() as temp_dir:
path = os.path.join(temp_dir, 'test')
original.save_to_file(path, password)
recovered = PasswordDataBase.load_from_file(path, password)
for expected, actual in zip(original.records, recovered.records):
self.assertEqual(expected, actual)
def test_load_fails_with_wrong_password(self):
"""
Ensure that records written to a file can't be recovered using an
incorrect password.
"""
database = PasswordDataBase()
database.records.append(
PasswordDataBase.Record(name='foo',
login='bar',
password=generate_password(64, ALL_CHARS)))
database.records.append(
PasswordDataBase.Record(name='ga',
login='bu',
password=generate_password(64, ALL_CHARS),
tags=['zo', 'meu']))
database.records.append(PasswordDataBase.Record())
with tempfile.TemporaryDirectory() as temp_dir:
path = os.path.join(temp_dir, 'test')
database.save_to_file(path, 'right')
self.assertRaises(InvalidPasswordException,
lambda: PasswordDataBase.load_from_file(path, 'wrong'))
if __name__ == '__main__':
unittest.main()
My main concerns are:
- Does this code follow best practices when it comes to security?
- Is the test coverage good enough? I feel like it could be improved, but I find it hard to come up with additional meaningful tests.
Any other remarks about my code are also welcome, of course.
