1
0

storage: replace STO_EV_* ints with IntEnum

This commit is contained in:
SomberNight
2019-09-04 13:31:49 +02:00
parent 54776ca1d9
commit a42a773d19
5 changed files with 32 additions and 29 deletions

View File

@@ -28,6 +28,7 @@ import stat
import hashlib
import base64
import zlib
from enum import IntEnum
from . import ecc
from .util import profiler, InvalidPassword, WalletFileException, bfh, standardize_path
@@ -42,9 +43,11 @@ def get_derivation_used_for_hw_device_encryption():
"/4541509'" # ascii 'ELE' as decimal ("BIP43 purpose")
"/1112098098'") # ascii 'BIE2' as decimal
# storage encryption version
STO_EV_PLAINTEXT, STO_EV_USER_PW, STO_EV_XPUB_PW = range(0, 3)
class StorageEncryptionVersion(IntEnum):
PLAINTEXT = 0
USER_PASSWORD = 1
XPUB_PASSWORD = 2
class WalletStorage(Logger):
@@ -67,7 +70,7 @@ class WalletStorage(Logger):
self.db = DB_Class(self.raw, manual_upgrades=manual_upgrades)
self.load_plugins()
else:
self._encryption_version = STO_EV_PLAINTEXT
self._encryption_version = StorageEncryptionVersion.PLAINTEXT
# avoid new wallets getting 'upgraded'
self.db = DB_Class('', manual_upgrades=False)
@@ -128,13 +131,13 @@ class WalletStorage(Logger):
def is_encrypted(self):
"""Return if storage encryption is currently enabled."""
return self.get_encryption_version() != STO_EV_PLAINTEXT
return self.get_encryption_version() != StorageEncryptionVersion.PLAINTEXT
def is_encrypted_with_user_pw(self):
return self.get_encryption_version() == STO_EV_USER_PW
return self.get_encryption_version() == StorageEncryptionVersion.USER_PASSWORD
def is_encrypted_with_hw_device(self):
return self.get_encryption_version() == STO_EV_XPUB_PW
return self.get_encryption_version() == StorageEncryptionVersion.XPUB_PASSWORD
def get_encryption_version(self):
"""Return the version of encryption used for this storage.
@@ -151,13 +154,13 @@ class WalletStorage(Logger):
try:
magic = base64.b64decode(self.raw)[0:4]
if magic == b'BIE1':
return STO_EV_USER_PW
return StorageEncryptionVersion.USER_PASSWORD
elif magic == b'BIE2':
return STO_EV_XPUB_PW
return StorageEncryptionVersion.XPUB_PASSWORD
else:
return STO_EV_PLAINTEXT
return StorageEncryptionVersion.PLAINTEXT
except:
return STO_EV_PLAINTEXT
return StorageEncryptionVersion.PLAINTEXT
@staticmethod
def get_eckey_from_password(password):
@@ -167,9 +170,9 @@ class WalletStorage(Logger):
def _get_encryption_magic(self):
v = self._encryption_version
if v == STO_EV_USER_PW:
if v == StorageEncryptionVersion.USER_PASSWORD:
return b'BIE1'
elif v == STO_EV_XPUB_PW:
elif v == StorageEncryptionVersion.XPUB_PASSWORD:
return b'BIE2'
else:
raise WalletFileException('no encryption magic for version: %s' % v)
@@ -211,13 +214,13 @@ class WalletStorage(Logger):
"""Set a password to be used for encrypting this storage."""
if enc_version is None:
enc_version = self._encryption_version
if password and enc_version != STO_EV_PLAINTEXT:
if password and enc_version != StorageEncryptionVersion.PLAINTEXT:
ec_key = self.get_eckey_from_password(password)
self.pubkey = ec_key.get_public_key_hex()
self._encryption_version = enc_version
else:
self.pubkey = None
self._encryption_version = STO_EV_PLAINTEXT
self._encryption_version = StorageEncryptionVersion.PLAINTEXT
# make sure next storage.write() saves changes
self.db.set_modified(True)