def test_update_new_token(self):
initial_count = Token.select().count()
count = {"token": "dummy", "frequency": 5}
Token.update_count(count, self.enc)
final_count = Token.select().count()
assert initial_count + 1 == final_count
entry = Token.get(Token.hash == self.enc.salted_hash("dummy"))
entry.frequency == 5
def test_batch_update(self):
initial_count = Token.select().count()
counts = [
{
"token": "batch-update-1",
"frequency": 5
},
{
"token": "batch-update-2",
"frequency": 3
},
]
Token.batch_update(counts, self.enc)
final_count = Token.select().count()
assert initial_count + len(counts) == final_count
def newtoken_v1():
data = json.loads(request.data.decode('utf8'))
print(data)
token_old = Token.query.filter_by(token=data['token']).one_or_none()
if token_old is None:
return json.dumps({"status": "ERROR", "error_msg": "Wrong token"})
tknstr = Config.SECRET_KEY + data['token']
token_new = Token(user_id=token_old.user_id)
token_new.generate(tknstr)
db.session.delete(token_old)
db.session.add(token_new)
db.session.commit()
return json.dumps({"status": "OK", "data": {"token": token_new.token}})
def test_decrypt(self):
count = {"token": "Sparta", "frequency": 5}
token = Token.update_count(count, self.enc)
decrypted = token.decrypt(self.enc)
assert token.encrypted != decrypted["token"]
assert self.enc.decrypt(token.encrypted) == decrypted["token"]
assert token.frequency == decrypted["frequency"]
def login_v1():
data = json.loads(request.data.decode('utf8'))
user = User.query.filter_by(username=data['login']).one_or_none()
if user is None:
return json.dumps({
"status": "ERROR",
"error_msg": "Login doesn't exist"
})
print(user.check_password(data['password']))
if not user.check_password(data['password']):
# need log it to analyze logs for ban bruteforcers
return json.dumps({"status": "ERROR", "error_msg": "Password wrong"})
tknstr = Config.SECRET_KEY + data['login'] + data['password']
token = Token(user_id=user.id)
token.generate(tknstr)
db.session.add(token)
db.session.commit()
res = {"id": str(user.id), "token": token.token}
return json.dumps({"status": "OK", "data": res})
def test_updating_existing_token(self):
Token.update_count({"token": "existing", "frequency": 2}, self.enc)
initial_count = Token.select().count()
count = {"token": "existing", "frequency": 5}
Token.update_count(count, self.enc)
final_count = Token.select().count()
assert initial_count == final_count
entry = Token.get(Token.hash == self.enc.salted_hash("existing"))
assert entry.frequency == 7
def test_get_page(self):
# Clear the table
Token.delete().execute()
words = ["Sparta", "Spartan", "Spartacus"]
counts = [{"token": word, "frequency": 1} for word in words]
Token.batch_update(counts, self.enc)
page1, total1 = Token.get_page(0, 1, self.enc)
page2, total2 = Token.get_page(1, 1, self.enc)
page3, total3 = Token.get_page(2, 1, self.enc)
# same page length
assert len(page1) == len(page2) == len(page3)
# content differs
assert page1 != page2 != page3
# total count is the same
assert total1 == total2 == total3
# Data is already decrypted
assert page1[0]["token"] in words