Advertisement
Not a member of Pastebin yet?
Sign Up,
it unlocks many cool features!
- 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()
Advertisement
Add Comment
Please, Sign In to add comment
Advertisement