As an exercise, I've decided to write a lightweight, dictionary type database. Below are some of the features I've implemented:
- Overwrite Inserting: I allow the user to determine if they want to overwrite existing data, should they insert a pair that already has a value associated with the key of that pair.
- Encryption: The user can pass a 32 character long password that encrypts the database when they're done using it.
- Query Specification: The user can request a value by passing a key, request all keys associated with a particular value, or pass a pair and get the index of where that pair is in the database.
Questions:
- Security: Is how I'm implementing security "good"? I've user
Fernetin the past, and the absolute requirement of a 32 character long password ensures it will take a considerable amount of time to break the encryption. I'm also fairly concerned about the time between each encrypt and decrypt. Should I only decrypt when the user wants to insert or query the database? - Password Verification: I let
Fernetdecide if the password is correct, instead of implementing something myself. Is this a good way of going about this? - Conventions: The ones I'm particular about are the double underscore ones, such as
__encrypt_db. I'm familiar with the purpose of hiding functions that are meant to be internal. Am I using this convention correctly? - Any other improvements that you think I can make.
lindb.py
"""
LinDB
@author Ben Antonellis.
@date 07-17-2020.
"""
import os
import json
import base64
from cryptography.fernet import Fernet
from cryptography.fernet import InvalidToken
from typing import Any, Union, List, Dict
class LinDB():
def __init__(self, name, pw=None):
self.name = name
self.__pw = pw
self.db = {}
self.file_name = f"{self.name}.json"
self.connected = False
self.new_db = False
self.encrypt = self.__pw != None
if self.encrypt:
if len(self.__pw) > 32:
raise PasswordLengthError("Password must be at least 32 characters long!")
self.__pw = base64.urlsafe_b64encode(self.__pw.encode())
self.fernet = Fernet(self.__pw)
self.__create_db_file()
def insert(self, pair: Dict, overwrite:bool=False) -> None:
"""
Allows the user to insert a dictionary into the database.
"""
if not self.connected:
quit("Please call .connect() to connect to database!")
for key in pair:
value = pair[key]
if overwrite:
for pair_key, db_key in zip(pair, self.db):
if pair_key == db_key:
self.db[db_key] = value
break
self.db.update(pair)
def query(self, key:Any=None, value:Any=None) -> Union[None, List[Any], bool]:
"""
Querys the database for either the key or value.
If both key and value:
Return position in database the first pair was found.
If just key:
Return value associated with key.
If just value:
Return all keys with associated value.
"""
if not self.connected:
quit("Please call .connect() to connect to database!")
try:
if key and value:
index = 0
for k, v in self.db.items():
if k == key and v == value:
return index
index += 1
if key and not value:
return self.db[key]
if value and not key:
return [k for k, v in self.db.items() if v == value]
except KeyError:
return
def save(self) -> None:
"""
Saves the current database to the file.
"""
if not self.connected:
quit("Please call .connect() to connect to database!")
with open(self.file_name, "w") as db_file:
json.dump(self.db, db_file, ensure_ascii=False)
def connect(self) -> None:
"""
Indicates to the database that it should start decrypting now.
"""
if self.__db_empty():
self.connected = True
return
if self.encrypt:
try:
if not self.new_db:
self.__decrypt_db()
self.connected = True
self.__load_db_file()
except InvalidToken:
quit("Wrong password for database!")
def done(self) -> None:
"""
Indicates to the database that it should start encrypting now.
"""
if not self.connected:
quit("Please call .connect() to connect to database!")
if self.encrypt:
self.__encrypt_db()
self.connected = False
def __create_db_file(self) -> None:
"""
Creates a database file with the name of the database as the filename.
"""
if not os.path.exists(self.file_name):
_ = open(self.file_name, "w").close()
self.new_db = True
def __load_db_file(self) -> None:
"""
Load the database into the current database dictionary.
"""
with open(self.file_name, "r") as db_file:
try:
json.load(db_file)
except json.decoder.JSONDecodeError:
print("Previous database not found. Creating new database.")
self.db = {}
def __encrypt_db(self) -> None:
"""
Encrypts the database with Fernet.
"""
with open(self.file_name, 'rb') as db_file:
db = db_file.readline()
encrypted = self.fernet.encrypt(db)
with open(self.file_name, 'wb') as db_file:
db_file.write(encrypted)
def __decrypt_db(self) -> None:
"""
Decrypts the database with Fernet.
"""
with open(self.file_name, 'rb') as db_file:
db = db_file.readline()
decrypted = self.fernet.decrypt(db)
with open(self.file_name, 'wb') as db_file:
db_file.write(decrypted)
def __db_empty(self) -> bool:
"""
Determines if the database if empty.
"""
with open(self.file_name, "r") as db_file:
return not db_file.readlines()
def __repr__(self):
return f"DB: {self.name}"
class PasswordLengthError(Exception):
"""
Raised when the user enters a password less than 32 characters long.
"""
def __init__(self, message):
super().__init__(message)
Below is an example file of how an average user would work with this database:
test_db.py
from lindb import LinDB
# Example password 32 characters long #
pw = "zSLfLhAvjhmX6CrzCbxSE2dzXEZaiOfO"
db = LinDB("DB_TEST", pw=pw)
# Decrypts the file if the password is correct #
db.connect()
# Start inserting pairs #
db.insert({"Ben": 16})
db.insert({"Hannah": 17})
db.insert({"Will": 18})
# Query database and display results #
results = [
db.query(value=16),
db.query(key="Hannah"),
db.query(key="Will", value=18),
db.query(key="Test")
]
for result in results:
print(result)
# Demonstrating the ability to use assignment expressions #
# Should the key and/or value not exist, None is returned #
if result := db.query(key="Be"):
print(result)
# This writes the current database to the file #
db.save()
# Encrypts the file #
db.done()
