From cb403a94bdcfa9b853273d297e2b4d3ffac23402 Mon Sep 17 00:00:00 2001 From: Inex Code Date: Fri, 30 Dec 2022 21:06:16 +0300 Subject: [PATCH] fix: typing --- selfprivacy_api/actions/api_tokens.py | 2 +- selfprivacy_api/jobs/__init__.py | 49 ++++++++++--------- .../tokens/abstract_tokens_repository.py | 11 +++-- .../tokens/json_tokens_repository.py | 8 +-- .../tokens/redis_tokens_repository.py | 46 ++++++++++------- 5 files changed, 67 insertions(+), 49 deletions(-) diff --git a/selfprivacy_api/actions/api_tokens.py b/selfprivacy_api/actions/api_tokens.py index 57828bc..38133fd 100644 --- a/selfprivacy_api/actions/api_tokens.py +++ b/selfprivacy_api/actions/api_tokens.py @@ -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. diff --git a/selfprivacy_api/jobs/__init__.py b/selfprivacy_api/jobs/__init__.py index 1547b84..fe4a053 100644 --- a/selfprivacy_api/jobs/__init__.py +++ b/selfprivacy_api/jobs/__init__.py @@ -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", diff --git a/selfprivacy_api/repositories/tokens/abstract_tokens_repository.py b/selfprivacy_api/repositories/tokens/abstract_tokens_repository.py index 931f64d..3a20ede 100644 --- a/selfprivacy_api/repositories/tokens/abstract_tokens_repository.py +++ b/selfprivacy_api/repositories/tokens/abstract_tokens_repository.py @@ -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") diff --git a/selfprivacy_api/repositories/tokens/json_tokens_repository.py b/selfprivacy_api/repositories/tokens/json_tokens_repository.py index 394c046..77e1311 100644 --- a/selfprivacy_api/repositories/tokens/json_tokens_repository.py +++ b/selfprivacy_api/repositories/tokens/json_tokens_repository.py @@ -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, } diff --git a/selfprivacy_api/repositories/tokens/redis_tokens_repository.py b/selfprivacy_api/repositories/tokens/redis_tokens_repository.py index d665553..c72e231 100644 --- a/selfprivacy_api/repositories/tokens/redis_tokens_repository.py +++ b/selfprivacy_api/repositories/tokens/redis_tokens_repository.py @@ -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))