async def otp() -> OtpData:
# Key is 12345
otp = OtpData(symptoms_started_on=date.today())
# Authorize this otp
await managers.otp_redis.set(
key_for_otp_sha(sha256("12345".encode("utf-8")).hexdigest()), OtpDataSchema().dumps(otp),
)
return otp
async def store(otp: str, otp_data: OtpData) -> None:
"""
Store the OtpData associated with the OTP, managing the key and value dump to the database.
:param otp: the OTP associated with the database entry.
:param otp_data: the OtpData to store.
:raises: OtpCollision if the OTP is already in the database.
"""
did_set = await managers.otp_redis.set(
key=_key_for_otp(otp),
value=OtpDataSchema().dumps(otp_data),
expire=config.OTP_KEY_EXPIRATION_SECONDS,
exist=StringCommandsMixin.SET_IF_NOT_EXIST,
)
if not did_set:
raise OtpCollisionException()
async def validate_otp_token(otp_sha: str, delete: bool = False) -> OtpData:
"""
Load an OtpData model from the database.
:param otp_sha: the sha256 of the OTP code.
:param delete: if true, deletes the key right after retrieving its content.
:return: the deserialized OtpData model associated with the given sha256 string.
:raises: UnauthorizedOtpException if there is no OtpData associated with the given sha256.
"""
key = key_for_otp_sha(otp_sha)
if not delete:
data = await managers.otp_redis.get(key)
else:
pipe = managers.otp_redis.pipeline()
pipe.get(key)
pipe.delete(key)
data = (await pipe.execute())[0]
if data is None:
raise UnauthorizedOtpException()
return OtpDataSchema().loads(data)
def test_otp_schema_fails_on_additional(invalid_json: Dict[str, str]) -> None:
with raises(ValidationError):
OtpDataSchema().load(invalid_json)
def test_otp_schema_fails_on_missing_required() -> None:
with raises(ValidationError):
OtpDataSchema().load(dict())
def test_otp_data_schema_success() -> None:
OtpDataSchema().load(_OTP_DATA_SERIALIZED)