ska-referee-test/ref-test/common/security/database.py
viveksantayana e1967bcd7e Finished data upload
Refactored to move security package inside common
Moved data folder to process root.
2021-11-28 02:30:46 +00:00

46 lines
1.8 KiB
Python

from pymongo import collection
from . import encrypt, decrypt
encrypted_parameters = ['username', 'email', 'name', 'club', 'creator']
def decrypt_find(collection:collection, query:dict):
cursor = collection.find({})
output_list = []
for document in cursor:
decrypted_document = {}
for key in document:
if key not in encrypted_parameters:
decrypted_document[key] = document[key]
else:
decrypted_document[key] = decrypt(document[key])
if not query:
output_list.append(decrypted_document)
else:
if set(query.items()).issubset(set(decrypted_document.items())):
output_list.append(decrypted_document)
return output_list
def decrypt_find_one(collection:collection, query:dict={}):
cursor = decrypt_find(collection=collection, query=query)
if cursor: return cursor[0]
return None
def encrypted_update(collection:collection, query:dict={}, update:dict={}):
document = decrypt_find_one(collection=collection, query=query)
for update_action in update:
key_pairs = update[update_action]
if type(key_pairs) is not dict:
raise ValueError
if update_action == '$set':
for key in key_pairs:
if key == '_id':
raise ValueError
document[key] = key_pairs[key]
if update_action == '$unset':
for key in key_pairs:
if key == '_id':
raise ValueError
if key in document:
del document[key]
for key in document:
document[key] = encrypt(document[key]) if key in encrypted_parameters else document[key]
return collection.find_one_and_replace( { '_id': document['_id'] }, document)