diff --git a/johnnycanencrypt/__init__.py b/johnnycanencrypt/__init__.py index 5a26f95..43d596f 100644 --- a/johnnycanencrypt/__init__.py +++ b/johnnycanencrypt/__init__.py @@ -4,6 +4,7 @@ import os import sqlite3 import urllib.parse +import shutil from datetime import datetime from enum import Enum from typing import Dict, List, Optional, Union, Tuple, Any @@ -183,6 +184,17 @@ def __str__(self) -> str: def upgrade_if_required(self): "Upgrades the database schema if required" + oldpath = self._upgrade_if_required() + if oldpath is None: + return + os.unlink(oldpath) + # Now let us rename the file + shutil.copy(self.dbpath, oldpath) + os.unlink(self.dbpath) + self.dbpath = oldpath + + def _upgrade_if_required(self): + "Internal: Upgrades the database schema if required" SHOULD_WE = False existing_records = [] con = sqlite3.connect(self.dbpath) @@ -205,6 +217,7 @@ def upgrade_if_required(self): existing_records = cursor.fetchall() else: return + con.close() # Temporay db setup oldpath = self.dbpath self.dbpath = self.path / "jce_upgrade.db" @@ -222,6 +235,7 @@ def upgrade_if_required(self): cursor.execute( "INSERT INTO dbupgrade (upgradedate) values (?)", (DB_UPGRADE_DATE,) ) + con.close() # now let us insert our existing data for row in existing_records: ( @@ -255,9 +269,8 @@ def upgrade_if_required(self): fingerprint = row["fingerprint"] sql = "UPDATE keys set oncard=?, primary_on_card=? where fingerprint=?" cursor.execute(sql, (oncard, primary_on_card, fingerprint)) - # Now let us rename the file - os.rename(self.dbpath, oldpath) - self.dbpath = oldpath + con.close() + return oldpath def update_password(self, key: Key, password: str, newpassword: str) -> Key: """Updates the password of the given key and saves to the database""" @@ -464,6 +477,7 @@ def _save_key_info_to_db( value = uid[uid_keyname] sql = f"INSERT INTO {tablename} (value, key_id, value_id) values (?, ?, ?)" cursor.execute(sql, (value, key_id, value_id)) + con.close() def __contains__(self, other: Union[str, Key]) -> bool: """Checks if a Key object of fingerprint str exists in the keystore or not. @@ -536,7 +550,7 @@ def update_expiry_in_subkeys( sql, (etime_str, subkey[1]), ) - + con.close() # Regnerate the key object and return it return self.get_key(fingerprint) @@ -605,7 +619,7 @@ def add_userid(self, key: Key, userid: str, password: str) -> Key: value = uid[uid_keyname] sql = f"INSERT INTO {tablename} (value, key_id, value_id) values (?, ?, ?)" cursor.execute(sql, (value, key_id, value_id)) - + con.close() # Regnerate the key object and return it return self.get_key(fingerprint) @@ -658,7 +672,7 @@ def revoke_userid(self, key: Key, userid: str, password: str) -> Key: revoked = 1 sql = "UPDATE uidvalues set revoked=? where id=?" cursor.execute(sql, (revoked, value_id)) - + con.close() # Regnerate the key object and return it return self.get_key(fingerprint) diff --git a/tests/test_sign_verify_bytes.py b/tests/test_sign_verify_bytes.py index bf20a83..f098afe 100644 --- a/tests/test_sign_verify_bytes.py +++ b/tests/test_sign_verify_bytes.py @@ -51,11 +51,11 @@ def test_sign_verify_file_cleartext(tmp_path): True, ) assert os.path.exists(output) - with open(output) as fobj: + with open(output, "rb") as fobj: data = fobj.read() - assert data.startswith("-----BEGIN PGP SIGNED MESSAGE-----") - assert "🦄🦄🦄" in data - assert data.endswith("-----END PGP SIGNATURE-----\n") + assert data.startswith(b"-----BEGIN PGP SIGNED MESSAGE-----") + assert "🦄🦄🦄".encode("utf-8") in data + assert data.endswith(b"-----END PGP SIGNATURE-----\n") jp = jce.Johnny(_get_cert_data(PUBLIC_PATH)) assert jp.verify_file(output.encode("utf-8")) diff --git a/tests/utils.py b/tests/utils.py index 331a7c9..6e4b8fb 100644 --- a/tests/utils.py +++ b/tests/utils.py @@ -6,9 +6,9 @@ def _get_cert_data(filepath): def verify_files(inputfile, decrypted_output): # read both the files - with open(inputfile) as f: + with open(inputfile, "rb") as f: original_text = f.read() - with open(decrypted_output) as f: + with open(decrypted_output, "rb") as f: decrypted_text = f.read() assert original_text == decrypted_text