dbmanager.py

import sqlite3
import hashlib
import base64
from tkinter import *

from cryptography.fernet import Fernet

DB_FILE_NAME = 'my.db'



class AuthenticationError(ValueError):
    pass
    

class EncryptedDB():

    class ToEncrypt(bytes):
        '''
        A wrapper class around bytes for passing to sqlite.
        '''
        
    @staticmethod
    def bytes_check(s, encoding='UTF-8'):
        if isinstance(s, str):
            s = bytes(s, encoding)
        if not isinstance(s, bytes):
            raise ValueError(
                'must be bytes or str, not {}'.format(type(s))
            )
        return s
    @staticmethod
    def derive_password(p):
        p = EncryptedDB.bytes_check(p)
        return hashlib.sha256(p).digest()
        
    
    
    def __init__(self):
        self.db_file = DB_FILE_NAME
        self.conn = sqlite3.connect(
                         self.db_file
                        ,detect_types=sqlite3.PARSE_DECLTYPES)
        self.conn.row_factory = sqlite3.Row
        self.conn.executescript('''
            CREATE TABLE IF NOT EXISTS version_info (
                version INTEGER
            );
            CREATE TABLE IF NOT EXISTS encryptedstring (
                id INTEGER PRIMARY KEY AUTOINCREMENT
               ,owner       TEXT
               ,data        ENCRYPTED
               ,description TEXT
               ,FOREIGN KEY(owner) REFERENCES login(username)
            );
            CREATE TABLE IF NOT EXISTS login (
                username TEXT PRIMARY KEY 
               ,password SHA256
               ,salt     BLOB
               ,userkey  ENCRYPTED
            );
        ''')

        self.fernet = None
        sqlite3.register_converter('SHA256', EncryptedDB.derive_password)

        
    def requires_login(f):
        def wrap(self, *args, **kwargs):
            if hasattr(self, 'logged_in_user'):
                r = f(self, *args, **kwargs)
            else:
                raise AuthenticationError(
                    'This method requires a prior call to login.'
                )
            return r
        return wrap
            
    
    def make_login(self, username, password):
        self.conn.execute('''
            insert into login(username, password) values(?,?)
        ''', (username, password))
        self.save()
        
    def login(self, user, password):
        r = self.conn.execute('''
            select * from login
            where username = ? and password = ?
        ''', (user, password)).fetchone()
        if r is None:
            raise AuthenticationError(
                'No such credential combination exists'
            )

        password = EncryptedDB.derive_password(password)
        password = base64.urlsafe_b64encode(password)
        self.fernet = Fernet(password)
        sqlite3.register_converter('ENCRYPTED', self.fernet.decrypt)
        sqlite3.register_adapter(EncryptedDB.ToEncrypt, self.fernet.encrypt)
        self.logged_in_user = user
    
    
    
    @requires_login
    def insert(self, data, description=None):
        data = EncryptedDB.bytes_check(data)
        data = EncryptedDB.ToEncrypt(data)
        if len(description) == 0:
            description = None
        self.conn.execute('''
            insert into encryptedstring(owner, data, description) values(?,?,?)
        ''', (self.logged_in_user, data, description) )
        
    @requires_login
    def get(self, *rowids):
        return self.conn.execute('''
                select * from encryptedstring
                where id in ({})
            '''.format(
                ','.join('?'*len(rowids))
                )
            , rowids).fetchall()
        
        
    @requires_login
    def list(self):
        return self.conn.execute('''
                select id, description
                from encryptedstring
                where owner = ?
                order by id
            ''', (self.logged_in_user,)).fetchall()
        
    def save(self):
        try:
            self.conn.execute('commit')
            return True
        except sqlite3.OperationalError:
            pass
        return False
        
    def __str__(self):
        if not hasattr(self, 'logged_in_user'):
            return repr(self)
        largest_id = self.conn.execute('''
                select id
                from encryptedstring
                
                order by id desc
            ''').fetchone()
        if largest_id is None:
            return ''
        largest_id = largest_id['id']
        fmt = f'{{:{largest_id}}}: {{}}\r'
        s = ''
        for r in self.list():
            s += fmt.format(*r)
        return s


class App(Tk):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.db = EncryptedDB()
        self.render_login()
    
    def render_login(self):
        self.num_failed_logins = 0
        f = Frame(self)  
        self.nentry = Entry(f, width=20)
        self.pentry = Entry(f, width=20, show=b'\xe2\x80\xa2'.decode())
        self.errorlabel = Label(self, fg='red')
        self.new_user = Button(self, text='New Login', command=self.add_creds)
        
        f.pack()
        self.nentry.grid(row=0,column=1)
        self.pentry.grid(row=1,column=1)
        Label(f, text='Username').grid(row=0,column=0)
        Label(f, text='Password').grid(row=1,column=0)
        Button(self, text='Submit', command=self.verify_creds).pack(anchor='center')
        Button(self, text='New Login', command=self.add_creds).pack(side=RIGHT,anchor='e')
        self.errorlabel.pack()
        
    def add_creds(self):
        u = self.nentry.get()
        p = self.pentry.get()
        if len(u) == 0 or len(p) == 0:
            self.errorlabel.config(text='cant be blank')
            return
        
        try:
            self.db.make_login(u, p)
        except sqlite3.IntegrityError:
            self.errorlabel.config(text='username taken')
            return
        
        self.verify_creds()

    
    def verify_creds(self):
        u = self.nentry.get()
        p = self.pentry.get()
        try:
            self.db.login(u,p)
        except AuthenticationError:
            if self.num_failed_logins == 0:
                self.errorlabel.config(text='invalid credentials')
            self.num_failed_logins += 1
            return
            
        # login correct, do whatever
        [w.pack_forget() for w in self.winfo_children()]
        self.render_main()

    def render_main(self):
        self.geometry('800x600')
        self.main_frame = Frame(self)
        self.main_frame.pack(expand=True, fill=BOTH)
        self.update_main_display()
        Button(self, text='New Entry', command=self.insert_popup).pack()
        
    def insert_popup(self):
        def do_insert():
            self.db.insert(
                data.get(1.0, END)
               ,description=desc.get()
            )
            self.db.save()
            self.update_main_display()
            w.destroy()
            
            
        w = Toplevel(self)
        f = Frame(w)
        f.pack()
        Label(f, text='Description:').pack()
        desc = Entry(f)
        desc.pack(expand=True, fill=X)
        Label(f, text='Secret Text:').pack()
        data = Text(f)
        data.pack(expand=True, fill=BOTH)
        Button(f, text='Submit', command=do_insert).pack()
        

    def update_main_display(self):
        for w in self.main_frame.winfo_children():
            w.pack_forget()
        for i, row in enumerate(self.db.list()):
            d = row['description']
            if d is None:
                d = '< No description >'
            Label(self.main_frame, text=str(i+1)).grid(row=i, column=0)
            l = Label(self.main_frame, text=d)
            l.id = row['id']
            l.bind('<Button-1>', lambda cb: self.show_secret(l.id))
            l.grid(row=i, column=1)
        
    def show_secret(self, id):
        s = self.db.get(id)
        # get returns a list but we are  just getting 1, so extract from list
        if len(s) == 1:
            s = s[0]['data'].decode()
        else:
            #this should never happen
            s = '< error: this should never happen >'
            raise RuntimeError(
                'record id from update_main_display returned '
                'more than one record. uh oh :)'
            )
            
        if s is not None:
            w = Toplevel(self)
            t = Text(w)
            t.insert(0.0, s)
            t.config(state=DISABLED)
            t.pack()
            


if __name__ == '__main__':
    App().mainloop()