You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
151 lines
4.4 KiB
151 lines
4.4 KiB
# python modules
|
|
import sys
|
|
import pickle
|
|
import base64
|
|
import getpass
|
|
import gzip
|
|
# own modules
|
|
from notebook import Notebook, Note
|
|
# cryptography module
|
|
from cryptography.fernet import Fernet, InvalidToken
|
|
from cryptography.hazmat.primitives import hashes
|
|
from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
|
|
from cryptography.hazmat.backends import default_backend
|
|
|
|
|
|
|
|
class Menu:
|
|
'''Display a menu and respond to choices when run.'''
|
|
|
|
def __init__(self):
|
|
self.salt = "a683c64de226677703f56e6b6ead94bbc3690ec5293c3de3ffdc"
|
|
self.savefile = 'notebook.safe'
|
|
self.notebook = Notebook()
|
|
self.choices = {
|
|
"1": self.show_notes,
|
|
"2": self.search_notes,
|
|
"3": self.add_note,
|
|
"4": self.modify_note,
|
|
"5": self.load_notes,
|
|
"6": self.save_notes,
|
|
"7": self.quit
|
|
}
|
|
|
|
def display_menu(self):
|
|
print("""
|
|
Notebook Menu
|
|
1. Show all Notes
|
|
2. Search Notes
|
|
3. Add Note
|
|
4. Modify Note
|
|
5. Load Notes
|
|
6. Save Notes
|
|
7. Quit """)
|
|
|
|
def run(self):
|
|
'''Display the menu and respond to choices.'''
|
|
while True:
|
|
self.display_menu()
|
|
choice = input("Enter an option: ")
|
|
action = self.choices.get(choice)
|
|
try:
|
|
action()
|
|
except TypeError:
|
|
print("{0} is not a valid choice".format(choice))
|
|
|
|
def show_notes(self, notes=None):
|
|
'''Display all notes stored in notebook object'''
|
|
if not notes:
|
|
notes = self.notebook.notes
|
|
for note in notes:
|
|
print("{0}: {1}\n{2}".format(note.id, note.tags, note.memo))
|
|
|
|
def search_notes(self):
|
|
'''Search for a note containing given string'''
|
|
filter = input("Search for: ")
|
|
notes = self.notebook.search(filter)
|
|
self.show_notes(notes)
|
|
|
|
def add_note(self):
|
|
'''Add a given not to notebook object'''
|
|
memo = input("Enter a memo: ")
|
|
self.notebook.new_note(memo)
|
|
print("Your note has been added.")
|
|
|
|
def modify_note(self):
|
|
'''Modify tag and memo of note with given id'''
|
|
id = input("Enter a note id: ")
|
|
memo = input("Enter a memo: ")
|
|
tags = input("Enter tags: ")
|
|
if memo:
|
|
if not self.notebook.modify_memo(id, memo):
|
|
print("Note with id {0} doesn't exist.".format(id))
|
|
return
|
|
if tags:
|
|
if not self.notebook.modify_tags(id, tags):
|
|
print("Note with id {0} doesn't exist.".format(id))
|
|
|
|
def load_notes(self):
|
|
'''Decrypt notebook safe file and load it into notebook object'''
|
|
try:
|
|
f = open(self.savefile, 'rb')
|
|
except IOError:
|
|
print("Could not open file")
|
|
else:
|
|
cipher = f.read()
|
|
f.close()
|
|
self.notebook = self._decode_notefile(cipher)
|
|
self.notebook._set_id()
|
|
|
|
def save_notes(self):
|
|
'''Encrypt notebook object and store it into notebook safe file'''
|
|
cipher = self._encode_notefile()
|
|
try:
|
|
f = open(self.savefile, 'wb')
|
|
except IOError:
|
|
print("Could not open file")
|
|
else:
|
|
f.write(cipher)
|
|
f.close()
|
|
|
|
def _decode_notefile(self, cipher):
|
|
crypt = Fernet(self._get_password())
|
|
try:
|
|
plain = crypt.decrypt(cipher)
|
|
except InvalidToken:
|
|
print("Wrong password")
|
|
else:
|
|
try:
|
|
plain = gzip.decompress(plain)
|
|
except OSError:
|
|
print("File not valid")
|
|
else:
|
|
return pickle.loads(plain)
|
|
|
|
def _encode_notefile(self):
|
|
plain = pickle.dumps(self.notebook, pickle.HIGHEST_PROTOCOL)
|
|
crypt = Fernet(self._get_password())
|
|
plain = gzip.compress(plain)
|
|
return crypt.encrypt(plain)
|
|
|
|
def _get_password(self):
|
|
'''Request passphrase and derive key from it'''
|
|
passphrase = getpass.getpass()
|
|
kdf = PBKDF2HMAC(
|
|
algorithm = hashes.SHA256(),
|
|
length = 32,
|
|
salt = self.salt.encode('utf-8'),
|
|
iterations = 10000,
|
|
backend = default_backend()
|
|
)
|
|
return base64.urlsafe_b64encode(kdf.derive(passphrase.encode('utf-8')))
|
|
|
|
def quit(self):
|
|
'''Quit application'''
|
|
print("Thank you for using your notebook today.")
|
|
sys.exit(0)
|
|
|
|
|
|
|
|
if __name__ == "__main__":
|
|
Menu().run()
|