learning-by-doing/python_3_oop/chapter02/menu.py

# 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()