fix: typing

redis/token-repo
Inex Code 2022-12-30 21:06:16 +03:00
parent 7b526b7bd9
commit cb403a94bd
5 changed files with 67 additions and 49 deletions

View File

@ -150,7 +150,7 @@ def get_new_device_auth_token() -> str:
return Mnemonic(language="english").to_mnemonic(bytes.fromhex(key.key))
def use_new_device_auth_token(mnemonic_phrase, name) -> str:
def use_new_device_auth_token(mnemonic_phrase, name) -> Optional[str]:
"""Use the new device auth token by converting the mnemonic string to a byte array.
If the mnemonic phrase is valid then generate a device token and return it.
New device auth token must be deleted.

View File

@ -97,8 +97,8 @@ class Jobs:
error=None,
result=None,
)
r = RedisPool().get_connection()
_store_job_as_hash(r, _redis_key_from_uuid(job.uid), job)
redis = RedisPool().get_connection()
_store_job_as_hash(redis, _redis_key_from_uuid(job.uid), job)
return job
@staticmethod
@ -113,10 +113,10 @@ class Jobs:
"""
Remove a job from the jobs list.
"""
r = RedisPool().get_connection()
redis = RedisPool().get_connection()
key = _redis_key_from_uuid(job_uuid)
if (r.exists(key)):
r.delete(key)
if redis.exists(key):
redis.delete(key)
return True
return False
@ -149,12 +149,12 @@ class Jobs:
if status in (JobStatus.FINISHED, JobStatus.ERROR):
job.finished_at = datetime.datetime.now()
r = RedisPool().get_connection()
redis = RedisPool().get_connection()
key = _redis_key_from_uuid(job.uid)
if r.exists(key):
_store_job_as_hash(r, key, job)
if redis.exists(key):
_store_job_as_hash(redis, key, job)
if status in (JobStatus.FINISHED, JobStatus.ERROR):
r.expire(key, JOB_EXPIRATION_SECONDS)
redis.expire(key, JOB_EXPIRATION_SECONDS)
return job
@ -163,10 +163,10 @@ class Jobs:
"""
Get a job from the jobs list.
"""
r = RedisPool().get_connection()
redis = RedisPool().get_connection()
key = _redis_key_from_uuid(uid)
if r.exists(key):
return _job_from_hash(r, key)
if redis.exists(key):
return _job_from_hash(redis, key)
return None
@staticmethod
@ -174,9 +174,14 @@ class Jobs:
"""
Get the jobs list.
"""
r = RedisPool().get_connection()
jobs = r.keys("jobs:*")
return [_job_from_hash(r, job_key) for job_key in jobs]
redis = RedisPool().get_connection()
job_keys = redis.keys("jobs:*")
jobs = []
for job_key in job_keys:
job = _job_from_hash(redis, job_key)
if job is not None:
jobs.append(job)
return jobs
@staticmethod
def is_busy() -> bool:
@ -189,11 +194,11 @@ class Jobs:
return False
def _redis_key_from_uuid(uuid):
return "jobs:" + str(uuid)
def _redis_key_from_uuid(uuid_string):
return "jobs:" + str(uuid_string)
def _store_job_as_hash(r, redis_key, model):
def _store_job_as_hash(redis, redis_key, model):
for key, value in model.dict().items():
if isinstance(value, uuid.UUID):
value = str(value)
@ -201,12 +206,12 @@ def _store_job_as_hash(r, redis_key, model):
value = value.isoformat()
if isinstance(value, JobStatus):
value = value.value
r.hset(redis_key, key, str(value))
redis.hset(redis_key, key, str(value))
def _job_from_hash(r, redis_key):
if r.exists(redis_key):
job_dict = r.hgetall(redis_key)
def _job_from_hash(redis, redis_key):
if redis.exists(redis_key):
job_dict = redis.hgetall(redis_key)
for date in [
"created_at",
"updated_at",

View File

@ -17,7 +17,7 @@ from selfprivacy_api.models.tokens.new_device_key import NewDeviceKey
class AbstractTokensRepository(ABC):
def get_token_by_token_string(self, token_string: str) -> Optional[Token]:
def get_token_by_token_string(self, token_string: str) -> Token:
"""Get the token by token"""
tokens = self.get_tokens()
for token in tokens:
@ -26,7 +26,7 @@ class AbstractTokensRepository(ABC):
raise TokenNotFound("Token not found!")
def get_token_by_name(self, token_name: str) -> Optional[Token]:
def get_token_by_name(self, token_name: str) -> Token:
"""Get the token by name"""
tokens = self.get_tokens()
for token in tokens:
@ -101,7 +101,12 @@ class AbstractTokensRepository(ABC):
if not self.is_recovery_key_valid():
raise RecoveryKeyNotFound("Recovery key not found")
recovery_hex_key = self.get_recovery_key().key
recovery_key = self.get_recovery_key()
if recovery_key is None:
raise RecoveryKeyNotFound("Recovery key not found")
recovery_hex_key = recovery_key.key
if not self._assert_mnemonic(recovery_hex_key, mnemonic_phrase):
raise RecoveryKeyNotFound("Recovery key not found")

View File

@ -85,13 +85,13 @@ class JsonTokensRepository(AbstractTokensRepository):
recovery_key = RecoveryKey.generate(expiration, uses_left)
with WriteUserData(UserDataFiles.TOKENS) as tokens_file:
expiration = recovery_key.expires_at
if expiration is not None:
expiration = expiration.strftime(DATETIME_FORMAT)
key_expiration: Optional[str] = None
if recovery_key.expires_at is not None:
key_expiration = recovery_key.expires_at.strftime(DATETIME_FORMAT)
tokens_file["recovery_token"] = {
"token": recovery_key.key,
"date": recovery_key.created_at.strftime(DATETIME_FORMAT),
"expiration": expiration,
"expiration": key_expiration,
"uses_left": recovery_key.uses_left,
}

View File

@ -32,29 +32,34 @@ class RedisTokensRepository(AbstractTokensRepository):
def get_tokens(self) -> list[Token]:
"""Get the tokens"""
r = self.connection
token_keys = r.keys(TOKENS_PREFIX + "*")
return [self._token_from_hash(key) for key in token_keys]
redis = self.connection
token_keys = redis.keys(TOKENS_PREFIX + "*")
tokens = []
for key in token_keys:
token = self._token_from_hash(key)
if token is not None:
tokens.append(token)
return tokens
def delete_token(self, input_token: Token) -> None:
"""Delete the token"""
r = self.connection
redis = self.connection
key = RedisTokensRepository._token_redis_key(input_token)
if input_token not in self.get_tokens():
raise TokenNotFound
r.delete(key)
redis.delete(key)
def reset(self):
for token in self.get_tokens():
self.delete_token(token)
self.delete_new_device_key()
r = self.connection
r.delete(RECOVERY_KEY_REDIS_KEY)
redis = self.connection
redis.delete(RECOVERY_KEY_REDIS_KEY)
def get_recovery_key(self) -> Optional[RecoveryKey]:
"""Get the recovery key"""
r = self.connection
if r.exists(RECOVERY_KEY_REDIS_KEY):
redis = self.connection
if redis.exists(RECOVERY_KEY_REDIS_KEY):
return self._recovery_key_from_hash(RECOVERY_KEY_REDIS_KEY)
return None
@ -74,8 +79,8 @@ class RedisTokensRepository(AbstractTokensRepository):
def delete_new_device_key(self) -> None:
"""Delete the new device key"""
r = self.connection
r.delete(NEW_DEVICE_KEY_REDIS_KEY)
redis = self.connection
redis.delete(NEW_DEVICE_KEY_REDIS_KEY)
@staticmethod
def _token_redis_key(token: Token) -> str:
@ -89,10 +94,13 @@ class RedisTokensRepository(AbstractTokensRepository):
def _decrement_recovery_token(self):
"""Decrement recovery key use count by one"""
if self.is_recovery_key_valid():
uses_left = self.get_recovery_key().uses_left
recovery_key = self.get_recovery_key()
if recovery_key is None:
return
uses_left = recovery_key.uses_left
if uses_left is not None:
r = self.connection
r.hset(RECOVERY_KEY_REDIS_KEY, "uses_left", uses_left - 1)
redis = self.connection
redis.hset(RECOVERY_KEY_REDIS_KEY, "uses_left", uses_left - 1)
def _get_stored_new_device_key(self) -> Optional[NewDeviceKey]:
"""Retrieves new device key that is already stored."""
@ -116,9 +124,9 @@ class RedisTokensRepository(AbstractTokensRepository):
d[key] = None
def _model_dict_from_hash(self, redis_key: str) -> Optional[dict]:
r = self.connection
if r.exists(redis_key):
token_dict = r.hgetall(redis_key)
redis = self.connection
if redis.exists(redis_key):
token_dict = redis.hgetall(redis_key)
RedisTokensRepository._prepare_model_dict(token_dict)
return token_dict
return None
@ -139,8 +147,8 @@ class RedisTokensRepository(AbstractTokensRepository):
return self._hash_as_model(redis_key, NewDeviceKey)
def _store_model_as_hash(self, redis_key, model):
r = self.connection
redis = self.connection
for key, value in model.dict().items():
if isinstance(value, datetime):
value = value.isoformat()
r.hset(redis_key, key, str(value))
redis.hset(redis_key, key, str(value))