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.remove_note,
            "6": self.load_notes,
            "7": self.save_notes,
            "8": self.__init__,
            "9": self.quit
        }

    def display_menu(self):
        print("""
Notebook Menu
1. Show all Notes
2. Search Notes
3. Add Note
4. Modify Note
5. Remove Note
6. Load Notes
7. Save Notes
8. Reset Notes
9. 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 remove_note(self):
        '''Remove note with given id from note list'''
        id = input("Enter a note id: ")
        if self.notebook._remove_note(id):
            print("Note with id {0} removed.".format(id))
        else:
            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()
            notebook = self._decode_notefile(cipher)
            if notebook:
                self.notebook = notebook
                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()