change password modularized

change-password.py

@@ -1,69 +1,27 @@
 #!/usr/bin/python
 
-import src.config_provider
-import src.sql
-import base64
 import getpass
-from Crypto.Cipher import AES
-from Crypto.Util import Counter
-import binascii
 
 import src.my_scrypt
+import src.sql
+import src.change_password
 
 config = src.config_provider.getConfig()
 src.sql.connect(config['Document']['documentPath'])
 
-currentPassword = getpass.getpass(prompt="Enter current password: ")
+current_password = getpass.getpass(prompt="Enter current password: ")
 
-currentPasswordHash = binascii.hexlify(src.my_scrypt.getVerificationHash(currentPassword))
+new_password1 = getpass.getpass(prompt="Enter new password: ")
+new_password2 = getpass.getpass(prompt="Repeat the same password: ")
 
-if currentPasswordHash != src.sql.getOption('password'):
-    print("Given password doesn't match hash")
-    exit(-1)
-
-currentPasswordEncryptionKey = src.my_scrypt.getEncryptionHash(currentPassword)
-
-newPassword1 = getpass.getpass(prompt="Enter new password: ")
-newPassword2 = getpass.getpass(prompt="Repeat the same password: ")
-
-if newPassword1 != newPassword2:
+if new_password1 != new_password2:
     print('Entered passwords are not identical!')
     exit(-1)
 
-newPasswordVerificationKey = binascii.hexlify(src.my_scrypt.getVerificationHash(newPassword1))
-newPasswordEncryptionKey = src.my_scrypt.getEncryptionHash(newPassword1)
+ret = src.change_password.change_password(current_password, new_password1)
 
-encryptedNotes = src.sql.getResults("select note_id, note_title, note_text from notes where encryption = 1")
-
-def decrypt(encryptedBase64):
-    encryptedBytes = base64.b64decode(encryptedBase64)
-
-    aes = getAes(currentPasswordEncryptionKey)
-    return aes.decrypt(encryptedBytes)
-
-def encrypt(plainText):
-    aes = getAes(newPasswordEncryptionKey)
-    encryptedBytes = aes.encrypt(plainText)
-
-    return base64.b64encode(encryptedBytes)
-
-def getAes(key):
-    return AES.new(key, AES.MODE_CTR, counter=Counter.new(128, initial_value=5))
-
-for note in encryptedNotes:
-    decryptedTitle = decrypt(note['note_title'])
-    decryptedText = decrypt(note['note_text'])
-
-    reEncryptedTitle = encrypt(decryptedTitle)
-    reEncryptedText = encrypt(decryptedText)
-
-    src.sql.execute("update notes set note_title = ?, note_text = ? where note_id = ?",
-                    [reEncryptedTitle, reEncryptedText, note['note_id']])
-
-    print("Note " + note['note_id'] + " re-encrypted with new password")
-
-src.sql.setOption('password', newPasswordVerificationKey)
-src.sql.commit()
-
-print("Changes committed. All encrypted notes were re-encrypted successfully with new password key.")
-print("You can now start application and login with new password.")
+if (ret['success']):
+    print("Changes committed. All encrypted notes were re-encrypted successfully with new password key.")
+    print("You can now start application and login with new password.")
+else:
+    print(ret['message'])

src/change_password.py

@@ -0,0 +1,57 @@
+import src.config_provider
+import src.sql
+import base64
+from Crypto.Cipher import AES
+from Crypto.Util import Counter
+import binascii
+
+import src.my_scrypt
+
+
+def change_password(current_password, new_password):
+    current_password_hash = binascii.hexlify(src.my_scrypt.getVerificationHash(current_password))
+
+    if current_password_hash != src.sql.getOption('password'):
+        return {
+            'success': False,
+            'message': "Given current password doesn't match hash"
+        }
+
+    current_password_encryption_key = src.my_scrypt.getEncryptionHash(current_password)
+
+    new_password_verification_key = binascii.hexlify(src.my_scrypt.getVerificationHash(new_password))
+    new_password_encryption_key = src.my_scrypt.getEncryptionHash(new_password)
+
+    encrypted_notes = src.sql.getResults("select note_id, note_title, note_text from notes where encryption = 1")
+
+    def decrypt(encrypted_base64):
+        encrypted_bytes = base64.b64decode(encrypted_base64)
+
+        aes = get_aes(current_password_encryption_key)
+        return aes.decrypt(encrypted_bytes)
+
+    def encrypt(plain_text):
+        aes = get_aes(new_password_encryption_key)
+        encryptedBytes = aes.encrypt(plain_text)
+
+        return base64.b64encode(encryptedBytes)
+
+    def get_aes(key):
+        return AES.new(key, AES.MODE_CTR, counter=Counter.new(128, initial_value=5))
+
+    for note in encrypted_notes:
+        decrypted_title = decrypt(note['note_title'])
+        decrypted_text = decrypt(note['note_text'])
+
+        re_encrypted_title = encrypt(decrypted_title)
+        re_encrypted_text = encrypt(decrypted_text)
+
+        src.sql.execute("update notes set note_title = ?, note_text = ? where note_id = ?",
+                        [re_encrypted_title, re_encrypted_text, note['note_id']])
+
+    src.sql.setOption('password', new_password_verification_key)
+    src.sql.commit()
+
+    return {
+        'success': True
+    }