Update homeFinance.py

homeFinance.py

@@ -1,20 +1,16 @@
-import base64
+from crypto_funcs import new_db, decrypt_db, encrypt_db
 import os
-from cryptography.fernet import Fernet, InvalidToken
-from cryptography.hazmat.backends import default_backend
-from cryptography.hazmat.primitives import hashes
-from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
 from tkinter import ttk, scrolledtext
 from tkinter import *
 import sys
-import getpass
 import time
 import numpy as np
 import pandas as pd
 
 # check if creating a new database or using old one
 new = True
-dbpath = None  # These two lines are to remove pycharm flagging upcoming lines as a warning
+mainfile = None
+dbpath = None  # These three lines are to remove pycharm flagging upcoming lines as a warning
 if len(sys.argv) < 2:
     new = True
 elif len(sys.argv) == 2:
@@ -25,45 +21,15 @@
     quit(1)
 
 if new:
-    while True:
-        p = getpass.getpass(prompt='Set New Password: ').encode()
-        p2 = getpass.getpass(prompt='Confirm New Password: ').encode()
-        if p == p2:
-            break
-        print('Passwords do not match')
-    salt = os.urandom(16)  # cryptographic salt to protect against rainbow table attacks
-    dbpath = input('Path/Name for the new db: ')
-    # metadata = database created timestamp, total transactions included, categories:accounts:payees, last opened time
-    metadata = str(time.time()) + ',' + '0' + ',' + 'Misc:Cash:Misc' + ',' + str(time.time()) + '\n'
-    kdf = PBKDF2HMAC(algorithm=hashes.SHA256(), length=32, salt=salt, iterations=100000, backend=default_backend())
-    key = base64.urlsafe_b64encode(kdf.derive(p))  # remove problematic characters
-    f_obj = Fernet(key)  # symmetric encyption
-    metadata_enc = f_obj.encrypt(metadata.encode())
-    with open(dbpath, 'wb') as f:
-        f.write(salt)
-        f.write(metadata_enc)
-    print('File saved as ' + dbpath)
+    new_db()
     print('Restart to begin')
     quit(0)
 else:
-    while True:
-        with open(dbpath, 'rb') as file_enc:
-            salt = file_enc.read(16)
-            p = getpass.getpass().encode()
-            kdf = PBKDF2HMAC(algorithm=hashes.SHA256(), length=32, salt=salt, iterations=100000,
-                             backend=default_backend())
-            key = base64.urlsafe_b64encode(kdf.derive(p))
-            f_obj = Fernet(key)
-            mainfile_enc = file_enc.read()
-            try:
-                mainfile = f_obj.decrypt(mainfile_enc)
-                print('Success.')
-                break
-            except InvalidToken:
-                print('Failed to decrypt. Re-enter password.')
-
-# parsing the data
+    mainfile = decrypt_db(dbpath)
+
+# parsing the data into a pandas dataframe
 metadata = mainfile.decode().split('\n')[0].split(',')
+# printing basic stats
 print('Database created at: ' + time.asctime(time.localtime(float(metadata[0]))))
 print('Database last opened at: ' + time.asctime(time.localtime(float(metadata[3]))))
 print('Total Transactions: ' + metadata[1])
@@ -668,24 +634,5 @@ def numeric_analyze(event=None):  # analyze and summarize the data numericaly
 root.mainloop()
 
 if input('Save (y/n): ').lower() == 'y':
-    while True:
-        pf = getpass.getpass(prompt='Enter Password: ').encode()
-        if pf == p:
-            break
-        print('Incorrect Password.')
-    salt = os.urandom(16)
-    metadata = metadata[0] + ',' + str(t_count) + ',' + ':'.join(
-        [';'.join(categories), ';'.join(accounts), ';'.join(payees)]) + ',' + str(time.time()) + '\n'
-    kdf = PBKDF2HMAC(algorithm=hashes.SHA256(), length=32, salt=salt, iterations=100000, backend=default_backend())
-    key = base64.urlsafe_b64encode(kdf.derive(pf))
-    f_obj = Fernet(key)
-    df_string = ''
-    for i, j in df.iterrows():
-        df_string = df_string + ','.join([str(x) for x in j]) + '\n'
-    df_string = df_string.rstrip()
-    final_enc = f_obj.encrypt(metadata.encode() + df_string.encode())
-    with open(dbpath, 'wb') as f:
-        f.write(salt)
-        f.write(final_enc)
-    print('File saved as ' + dbpath)
+    encrypt_db(dbpath, p, metadata[0], df, t_count, categories, accounts, payees)
     quit(0)