Major refactor of the code
A lot of function are now using classes instead of parameters or tuples isort was added to the dev requirements Comments were adjusted
This commit is contained in:
parent
bbc9650357
commit
3f744723a9
@ -1,2 +1,3 @@
|
|||||||
black
|
black
|
||||||
flake8
|
flake8
|
||||||
|
isort
|
||||||
|
@ -6,8 +6,9 @@ from . import (
|
|||||||
account_checks,
|
account_checks,
|
||||||
account_parsing,
|
account_parsing,
|
||||||
bot,
|
bot,
|
||||||
encryption,
|
classes,
|
||||||
database,
|
database,
|
||||||
|
encryption,
|
||||||
generate_password,
|
generate_password,
|
||||||
)
|
)
|
||||||
|
|
||||||
@ -15,6 +16,7 @@ __all__ = [
|
|||||||
"account_checks",
|
"account_checks",
|
||||||
"account_parsing",
|
"account_parsing",
|
||||||
"bot",
|
"bot",
|
||||||
|
"classes",
|
||||||
"encryption",
|
"encryption",
|
||||||
"database",
|
"database",
|
||||||
"generate_password",
|
"generate_password",
|
||||||
|
@ -1,5 +1,7 @@
|
|||||||
import string
|
import string
|
||||||
|
|
||||||
|
from .classes import DecryptedAccount
|
||||||
|
|
||||||
FORBIDDEN_CHARS = frozenset("`\n")
|
FORBIDDEN_CHARS = frozenset("`\n")
|
||||||
PUNCTUATION = frozenset(string.punctuation).difference(FORBIDDEN_CHARS)
|
PUNCTUATION = frozenset(string.punctuation).difference(FORBIDDEN_CHARS)
|
||||||
|
|
||||||
@ -24,13 +26,13 @@ def check_password(passwd: str) -> bool:
|
|||||||
return _base_check(passwd)
|
return _base_check(passwd)
|
||||||
|
|
||||||
|
|
||||||
def check_account(name: str, login: str, passwd: str) -> bool:
|
def check_account(account: DecryptedAccount) -> bool:
|
||||||
"""Runs checks for account name, login and password"""
|
"""Runs checks for account name, login and password"""
|
||||||
return all(
|
return all(
|
||||||
(
|
(
|
||||||
check_account_name(name),
|
check_account_name(account.name),
|
||||||
check_login(login),
|
check_login(account.login),
|
||||||
check_password(passwd),
|
check_password(account.password),
|
||||||
)
|
)
|
||||||
)
|
)
|
||||||
|
|
||||||
|
@ -1,37 +1,50 @@
|
|||||||
import io
|
import io
|
||||||
from typing import Iterator, Self, Type
|
from typing import Iterable, Self
|
||||||
|
|
||||||
import pydantic
|
import pydantic
|
||||||
|
|
||||||
|
from .classes import DecryptedAccount
|
||||||
|
|
||||||
|
|
||||||
class _Account(pydantic.BaseModel):
|
class _Account(pydantic.BaseModel):
|
||||||
name: str
|
name: str
|
||||||
login: str
|
login: str
|
||||||
password: str
|
password: str
|
||||||
|
|
||||||
@classmethod
|
def to_usual_account(self, user_id: int) -> DecryptedAccount:
|
||||||
def from_tuple(cls: Type[Self], tuple_: tuple[str, str, str]) -> Self:
|
return DecryptedAccount(
|
||||||
return cls(name=tuple_[0], login=tuple_[1], password=tuple_[2])
|
user_id=user_id,
|
||||||
|
name=self.name,
|
||||||
|
login=self.login,
|
||||||
|
password=self.password,
|
||||||
|
)
|
||||||
|
|
||||||
def as_tuple(self: Self) -> tuple[str, str, str]:
|
@classmethod
|
||||||
return (self.name, self.login, self.password)
|
def from_usual_account(cls, account: DecryptedAccount) -> Self:
|
||||||
|
return cls(
|
||||||
|
name=account.name,
|
||||||
|
login=account.login,
|
||||||
|
password=account.password,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
class _Accounts(pydantic.BaseModel):
|
class _Accounts(pydantic.BaseModel):
|
||||||
accounts: list[_Account] = pydantic.Field(default_factory=list)
|
accounts: list[_Account] = pydantic.Field(default_factory=list)
|
||||||
|
|
||||||
|
|
||||||
def _accounts_list_to_json(accounts: Iterator[tuple[str, str, str]]) -> str:
|
def _accounts_list_to_json(accounts: Iterable[DecryptedAccount]) -> str:
|
||||||
accounts = _Accounts(accounts=[_Account.from_tuple(i) for i in accounts])
|
result = _Accounts(
|
||||||
return accounts.json(ensure_ascii=False)
|
accounts=[_Account.from_usual_account(i) for i in accounts],
|
||||||
|
).json(ensure_ascii=False)
|
||||||
|
return result
|
||||||
|
|
||||||
|
|
||||||
def json_to_accounts(json_: str) -> list[tuple[str, str, str]]:
|
def json_to_accounts(json_: str, user_id: int) -> list[DecryptedAccount]:
|
||||||
accounts = _Accounts.parse_raw(json_)
|
accounts = _Accounts.parse_raw(json_)
|
||||||
return [i.as_tuple() for i in accounts.accounts]
|
return [account.to_usual_account(user_id) for account in accounts.accounts]
|
||||||
|
|
||||||
|
|
||||||
def accounts_to_json(accounts: Iterator[tuple[str, str, str]]) -> io.StringIO:
|
def accounts_to_json(accounts: Iterable[DecryptedAccount]) -> io.StringIO:
|
||||||
file = io.StringIO(_accounts_list_to_json(accounts))
|
file = io.StringIO(_accounts_list_to_json(accounts))
|
||||||
file.name = "passwords.json"
|
file.name = "passwords.json"
|
||||||
return file
|
return file
|
||||||
|
@ -5,7 +5,7 @@ import time
|
|||||||
import telebot
|
import telebot
|
||||||
from sqlalchemy.future import Engine
|
from sqlalchemy.future import Engine
|
||||||
|
|
||||||
from .. import encryption, database, generate_password
|
from .. import database, encryption, generate_password
|
||||||
from ..account_checks import (
|
from ..account_checks import (
|
||||||
check_account,
|
check_account,
|
||||||
check_account_name,
|
check_account_name,
|
||||||
@ -13,6 +13,7 @@ from ..account_checks import (
|
|||||||
check_password,
|
check_password,
|
||||||
)
|
)
|
||||||
from ..account_parsing import accounts_to_json, json_to_accounts
|
from ..account_parsing import accounts_to_json, json_to_accounts
|
||||||
|
from ..classes import DecryptedAccount
|
||||||
|
|
||||||
Message = telebot.types.Message
|
Message = telebot.types.Message
|
||||||
|
|
||||||
@ -31,9 +32,9 @@ def _send_tmp_message(
|
|||||||
def _base_handler(
|
def _base_handler(
|
||||||
bot: telebot.TeleBot, mes: Message, prev_mes: Message | None = None
|
bot: telebot.TeleBot, mes: Message, prev_mes: Message | None = None
|
||||||
) -> None:
|
) -> None:
|
||||||
bot.delete_message(mes.chat.id, mes.id)
|
bot.delete_message(mes.chat.id, mes.id, timeout=5)
|
||||||
if prev_mes is not None:
|
if prev_mes is not None:
|
||||||
bot.delete_message(prev_mes.chat.id, prev_mes.id)
|
bot.delete_message(prev_mes.chat.id, prev_mes.id, timeout=5)
|
||||||
|
|
||||||
|
|
||||||
def get_accounts(bot: telebot.TeleBot, engine: Engine, mes: Message) -> None:
|
def get_accounts(bot: telebot.TeleBot, engine: Engine, mes: Message) -> None:
|
||||||
@ -68,11 +69,11 @@ def delete_all(bot: telebot.TeleBot, engine: Engine, mes: Message) -> None:
|
|||||||
"Отправьте YES для подтверждения",
|
"Отправьте YES для подтверждения",
|
||||||
)
|
)
|
||||||
bot.register_next_step_handler(
|
bot.register_next_step_handler(
|
||||||
mes, functools.partial(_delete_all, bot, engine, bot_mes)
|
mes, functools.partial(_delete_all2, bot, engine, bot_mes)
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
def _delete_all(
|
def _delete_all2(
|
||||||
bot: telebot.TeleBot, engine: Engine, prev_mes: Message, mes: Message
|
bot: telebot.TeleBot, engine: Engine, prev_mes: Message, mes: Message
|
||||||
) -> None:
|
) -> None:
|
||||||
_base_handler(bot, mes, prev_mes)
|
_base_handler(bot, mes, prev_mes)
|
||||||
@ -115,15 +116,12 @@ def _set_master_pass2(
|
|||||||
if text == "/cancel":
|
if text == "/cancel":
|
||||||
return _send_tmp_message(bot, mes.chat.id, "Успешная отмена")
|
return _send_tmp_message(bot, mes.chat.id, "Успешная отмена")
|
||||||
|
|
||||||
password_hash, master_salt = encryption.master_pass.encrypt_master_pass(
|
master_password = encryption.master_pass.encrypt_master_pass(
|
||||||
|
mes.from_user.id,
|
||||||
text,
|
text,
|
||||||
)
|
)
|
||||||
database.add.add_master_pass(
|
database.add.add_master_pass(engine, master_password)
|
||||||
engine,
|
|
||||||
mes.from_user.id,
|
|
||||||
master_salt,
|
|
||||||
password_hash,
|
|
||||||
)
|
|
||||||
_send_tmp_message(bot, mes.chat.id, "Успех")
|
_send_tmp_message(bot, mes.chat.id, "Успех")
|
||||||
del mes, text
|
del mes, text
|
||||||
gc.collect()
|
gc.collect()
|
||||||
@ -135,13 +133,16 @@ def reset_master_pass(
|
|||||||
mes: Message,
|
mes: Message,
|
||||||
) -> None:
|
) -> None:
|
||||||
_base_handler(bot, mes)
|
_base_handler(bot, mes)
|
||||||
|
|
||||||
if database.get.get_master_pass(engine, mes.from_user.id) is None:
|
if database.get.get_master_pass(engine, mes.from_user.id) is None:
|
||||||
return _send_tmp_message(bot, mes.chat.id, "Мастер пароль не задан")
|
return _send_tmp_message(bot, mes.chat.id, "Мастер пароль не задан")
|
||||||
|
|
||||||
bot_mes = bot.send_message(
|
bot_mes = bot.send_message(
|
||||||
mes.chat.id,
|
mes.chat.id,
|
||||||
"Отправьте новый мастер пароль, осторожно, все текущие аккаунты "
|
"Отправьте новый мастер пароль, осторожно, все текущие аккаунты "
|
||||||
"будут удалены навсегда",
|
"будут удалены навсегда",
|
||||||
)
|
)
|
||||||
|
|
||||||
bot.register_next_step_handler(
|
bot.register_next_step_handler(
|
||||||
mes, functools.partial(_reset_master_pass2, bot, engine, bot_mes)
|
mes, functools.partial(_reset_master_pass2, bot, engine, bot_mes)
|
||||||
)
|
)
|
||||||
@ -155,9 +156,13 @@ def _reset_master_pass2(
|
|||||||
if text == "/cancel":
|
if text == "/cancel":
|
||||||
return _send_tmp_message(bot, mes.chat.id, "Успешная отмена")
|
return _send_tmp_message(bot, mes.chat.id, "Успешная отмена")
|
||||||
|
|
||||||
hash_, salt = encryption.master_pass.encrypt_master_pass(text)
|
master_password = encryption.master_pass.encrypt_master_pass(
|
||||||
|
mes.from_user.id,
|
||||||
|
text,
|
||||||
|
)
|
||||||
database.delete.purge_accounts(engine, mes.from_user.id)
|
database.delete.purge_accounts(engine, mes.from_user.id)
|
||||||
database.change.change_master_pass(engine, mes.from_user.id, salt, hash_)
|
database.change.change_master_pass(engine, master_password)
|
||||||
|
|
||||||
_send_tmp_message(
|
_send_tmp_message(
|
||||||
bot, mes.chat.id, "Все ваши аккаунты удалены, а мастер пароль изменён"
|
bot, mes.chat.id, "Все ваши аккаунты удалены, а мастер пароль изменён"
|
||||||
)
|
)
|
||||||
@ -267,24 +272,30 @@ def _add_account5(
|
|||||||
if text == "/cancel":
|
if text == "/cancel":
|
||||||
return _send_tmp_message(bot, mes.chat.id, "Успешная отмена")
|
return _send_tmp_message(bot, mes.chat.id, "Успешная отмена")
|
||||||
|
|
||||||
salt, hash_ = database.get.get_master_pass(engine, mes.from_user.id)
|
master_password = database.get.get_master_pass(engine, mes.from_user.id)
|
||||||
if not encryption.master_pass.check_master_pass(text, hash_, salt):
|
if not encryption.master_pass.check_master_pass(text, master_password):
|
||||||
return _send_tmp_message(
|
return _send_tmp_message(
|
||||||
bot,
|
bot,
|
||||||
mes.chat.id,
|
mes.chat.id,
|
||||||
"Не подходит главный пароль",
|
"Не подходит главный пароль",
|
||||||
)
|
)
|
||||||
|
|
||||||
name, login, passwd = data["name"], data["login"], data["passwd"]
|
# name, login, passwd = data["name"], data["login"], data["passwd"]
|
||||||
|
account = DecryptedAccount(
|
||||||
|
user_id=mes.from_user.id,
|
||||||
|
name=data["name"],
|
||||||
|
login=data["login"],
|
||||||
|
password=data["passwd"],
|
||||||
|
)
|
||||||
|
|
||||||
enc_login, enc_pass, salt = encryption.other_accounts.encrypt(
|
encrypted_account = encryption.other_accounts.encrypt(
|
||||||
login,
|
account,
|
||||||
passwd,
|
|
||||||
text,
|
text,
|
||||||
)
|
)
|
||||||
|
|
||||||
result = database.add.add_account(
|
result = database.add.add_account(
|
||||||
engine, mes.from_user.id, name, salt, enc_login, enc_pass
|
engine,
|
||||||
|
encrypted_account,
|
||||||
)
|
)
|
||||||
|
|
||||||
_send_tmp_message(
|
_send_tmp_message(
|
||||||
@ -293,19 +304,19 @@ def _add_account5(
|
|||||||
"Успех" if result else "Произошла не предвиденная ошибка",
|
"Успех" if result else "Произошла не предвиденная ошибка",
|
||||||
)
|
)
|
||||||
|
|
||||||
del data, name, login, passwd, enc_login
|
del data, account
|
||||||
|
|
||||||
gc.collect()
|
gc.collect()
|
||||||
|
|
||||||
|
|
||||||
def get_account(bot: telebot.TeleBot, engine: Engine, mes: Message) -> None:
|
def get_account(bot: telebot.TeleBot, engine: Engine, mes: Message) -> None:
|
||||||
_base_handler(bot, mes)
|
_base_handler(bot, mes)
|
||||||
bot_mes = bot.send_message(mes.chat.id, "Отправьте название аккаунта")
|
|
||||||
|
|
||||||
master_pass = database.get.get_master_pass(engine, mes.from_user.id)
|
master_pass = database.get.get_master_pass(engine, mes.from_user.id)
|
||||||
if master_pass is None:
|
if master_pass is None:
|
||||||
return _send_tmp_message(bot, mes.chat.id, "Нет мастер пароля")
|
return _send_tmp_message(bot, mes.chat.id, "Нет мастер пароля")
|
||||||
|
|
||||||
|
bot_mes = bot.send_message(mes.chat.id, "Отправьте название аккаунта")
|
||||||
bot.register_next_step_handler(
|
bot.register_next_step_handler(
|
||||||
mes, functools.partial(_get_account2, bot, engine, bot_mes)
|
mes, functools.partial(_get_account2, bot, engine, bot_mes)
|
||||||
)
|
)
|
||||||
@ -340,36 +351,28 @@ def _get_account3(
|
|||||||
if text == "/cancel":
|
if text == "/cancel":
|
||||||
return _send_tmp_message(bot, mes.chat.id, "Успешная отмена")
|
return _send_tmp_message(bot, mes.chat.id, "Успешная отмена")
|
||||||
|
|
||||||
master_salt, hash_pass = database.get.get_master_pass(
|
master_password = database.get.get_master_pass(
|
||||||
engine,
|
engine,
|
||||||
mes.from_user.id,
|
mes.from_user.id,
|
||||||
)
|
)
|
||||||
|
|
||||||
if not encryption.master_pass.check_master_pass(
|
if not encryption.master_pass.check_master_pass(text, master_password):
|
||||||
text,
|
|
||||||
hash_pass,
|
|
||||||
master_salt,
|
|
||||||
):
|
|
||||||
return _send_tmp_message(bot, mes.chat.id, "Не подходит мастер пароль")
|
return _send_tmp_message(bot, mes.chat.id, "Не подходит мастер пароль")
|
||||||
|
|
||||||
salt, enc_login, enc_pass = database.get.get_account_info(
|
account = database.get.get_account_info(engine, mes.from_user.id, name)
|
||||||
engine, mes.from_user.id, name
|
account = encryption.other_accounts.decrypt(
|
||||||
)
|
account,
|
||||||
login, passwd = encryption.other_accounts.decrypt(
|
|
||||||
enc_login,
|
|
||||||
enc_pass,
|
|
||||||
text,
|
text,
|
||||||
salt,
|
|
||||||
)
|
)
|
||||||
_send_tmp_message(
|
_send_tmp_message(
|
||||||
bot,
|
bot,
|
||||||
mes.chat.id,
|
mes.chat.id,
|
||||||
f"Логин:\n`{login}`\nПароль:\n`{passwd}`\nНажмите на логин "
|
f"Логин:\n`{account.login}`\nПароль:\n`{account.password}`\nНажмите "
|
||||||
"или пароль, чтобы скопировать",
|
"на логин или пароль, чтобы скопировать",
|
||||||
30,
|
30,
|
||||||
)
|
)
|
||||||
|
|
||||||
del text, mes, passwd, login
|
del text, mes
|
||||||
gc.collect()
|
gc.collect()
|
||||||
|
|
||||||
|
|
||||||
@ -455,15 +458,11 @@ def _export2(
|
|||||||
if text == "/cancel":
|
if text == "/cancel":
|
||||||
return _send_tmp_message(bot, mes.chat.id, "Успешная отмена")
|
return _send_tmp_message(bot, mes.chat.id, "Успешная отмена")
|
||||||
|
|
||||||
master_salt, hash_pass = database.get.get_master_pass(
|
master_password = database.get.get_master_pass(
|
||||||
engine,
|
engine,
|
||||||
mes.from_user.id,
|
mes.from_user.id,
|
||||||
)
|
)
|
||||||
if not encryption.master_pass.check_master_pass(
|
if not encryption.master_pass.check_master_pass(text, master_password):
|
||||||
text,
|
|
||||||
hash_pass,
|
|
||||||
master_salt,
|
|
||||||
):
|
|
||||||
return _send_tmp_message(bot, mes.chat.id, "Не подходит мастер пароль")
|
return _send_tmp_message(bot, mes.chat.id, "Не подходит мастер пароль")
|
||||||
|
|
||||||
accounts = database.get.get_all_accounts(engine, mes.from_user.id)
|
accounts = database.get.get_all_accounts(engine, mes.from_user.id)
|
||||||
@ -474,7 +473,7 @@ def _export2(
|
|||||||
del text, accounts, json_io
|
del text, accounts, json_io
|
||||||
gc.collect()
|
gc.collect()
|
||||||
time.sleep(30)
|
time.sleep(30)
|
||||||
bot.delete_message(bot_mes.chat.id, bot_mes.id)
|
bot.delete_message(bot_mes.chat.id, bot_mes.id, timeout=5)
|
||||||
|
|
||||||
|
|
||||||
def import_accounts(
|
def import_accounts(
|
||||||
@ -513,13 +512,16 @@ def _import2(
|
|||||||
mes.chat.id,
|
mes.chat.id,
|
||||||
"Вы должны отправить документ",
|
"Вы должны отправить документ",
|
||||||
)
|
)
|
||||||
if mes.document.file_size > 102_400: # If file size is bigger that 100 MB
|
if mes.document.file_size > 102_400: # If file size is bigger than 100 MB
|
||||||
return _send_tmp_message(bot, mes.chat.id, "Файл слишком большой")
|
return _send_tmp_message(bot, mes.chat.id, "Файл слишком большой")
|
||||||
|
|
||||||
file_info = bot.get_file(mes.document.file_id)
|
file_info = bot.get_file(mes.document.file_id)
|
||||||
downloaded_file = bot.download_file(file_info.file_path)
|
downloaded_file = bot.download_file(file_info.file_path)
|
||||||
try:
|
try:
|
||||||
accounts = json_to_accounts(downloaded_file.decode("utf-8"))
|
accounts = json_to_accounts(
|
||||||
|
downloaded_file.decode("utf-8"),
|
||||||
|
mes.from_user.id,
|
||||||
|
)
|
||||||
except Exception:
|
except Exception:
|
||||||
return _send_tmp_message(
|
return _send_tmp_message(
|
||||||
bot,
|
bot,
|
||||||
@ -537,7 +539,7 @@ def _import3(
|
|||||||
bot: telebot.TeleBot,
|
bot: telebot.TeleBot,
|
||||||
engine: Engine,
|
engine: Engine,
|
||||||
prev_mes: Message,
|
prev_mes: Message,
|
||||||
accounts: list[tuple[str, str, str]],
|
accounts: list[DecryptedAccount],
|
||||||
mes: Message,
|
mes: Message,
|
||||||
) -> None:
|
) -> None:
|
||||||
_base_handler(bot, mes, prev_mes)
|
_base_handler(bot, mes, prev_mes)
|
||||||
@ -545,40 +547,33 @@ def _import3(
|
|||||||
if text == "/cancel":
|
if text == "/cancel":
|
||||||
return _send_tmp_message(bot, mes.chat.id, "Успешная отмена")
|
return _send_tmp_message(bot, mes.chat.id, "Успешная отмена")
|
||||||
|
|
||||||
master_salt, hash_pass = database.get.get_master_pass(
|
master_password = database.get.get_master_pass(
|
||||||
engine,
|
engine,
|
||||||
mes.from_user.id,
|
mes.from_user.id,
|
||||||
)
|
)
|
||||||
if not encryption.master_pass.check_master_pass(
|
if not encryption.master_pass.check_master_pass(text, master_password):
|
||||||
text,
|
|
||||||
hash_pass,
|
|
||||||
master_salt,
|
|
||||||
):
|
|
||||||
return _send_tmp_message(bot, mes.chat.id, "Не подходит мастер пароль")
|
return _send_tmp_message(bot, mes.chat.id, "Не подходит мастер пароль")
|
||||||
|
|
||||||
# List of names of accounts, which failed to be added to the database
|
# List of names of accounts, which failed to be added to the database
|
||||||
# or failed the tests
|
# or failed the tests
|
||||||
failed: list[str] = []
|
failed: list[str] = []
|
||||||
for account in accounts:
|
for account in accounts:
|
||||||
name, login, passwd = account
|
if not check_account(account):
|
||||||
if not check_account(name, login, passwd):
|
failed.append(account.name)
|
||||||
failed.append(name)
|
|
||||||
continue
|
continue
|
||||||
enc_login, enc_passwd, salt = encryption.other_accounts.encrypt(
|
account = encryption.other_accounts.encrypt(
|
||||||
login,
|
account,
|
||||||
passwd,
|
|
||||||
text,
|
text,
|
||||||
)
|
)
|
||||||
result = database.add.add_account(
|
result = database.add.add_account(engine, account)
|
||||||
engine, mes.from_user.id, name, salt, enc_login, enc_passwd
|
|
||||||
)
|
|
||||||
if not result:
|
if not result:
|
||||||
failed.append(name)
|
failed.append(account.name)
|
||||||
|
|
||||||
if failed:
|
if failed:
|
||||||
mes_text = "Не удалось добавить:\n" + "\n".join(failed)
|
mes_text = "Не удалось добавить:\n" + "\n".join(failed)
|
||||||
else:
|
else:
|
||||||
mes_text = "Успех"
|
mes_text = "Успех"
|
||||||
|
|
||||||
_send_tmp_message(bot, mes.chat.id, mes_text, 10)
|
_send_tmp_message(bot, mes.chat.id, mes_text, 10)
|
||||||
del text, mes, accounts
|
del text, mes, accounts
|
||||||
gc.collect()
|
gc.collect()
|
||||||
|
21
src/classes.py
Normal file
21
src/classes.py
Normal file
@ -0,0 +1,21 @@
|
|||||||
|
from typing import Self, TypeAlias
|
||||||
|
|
||||||
|
import pydantic
|
||||||
|
|
||||||
|
from .database import models
|
||||||
|
|
||||||
|
Account: TypeAlias = models.Account
|
||||||
|
|
||||||
|
|
||||||
|
class DecryptedAccount(pydantic.BaseModel):
|
||||||
|
user_id: int
|
||||||
|
name: str
|
||||||
|
login: str
|
||||||
|
password: str
|
||||||
|
|
||||||
|
@classmethod
|
||||||
|
def from_tuple(cls, tuple_: tuple[str, str, str]) -> Self:
|
||||||
|
return cls(name=tuple_[0], login=tuple_[1], password=tuple_[2])
|
||||||
|
|
||||||
|
def as_tuple(self) -> tuple[str, str, str]:
|
||||||
|
return (self.name, self.login, self.password)
|
@ -1,3 +1,3 @@
|
|||||||
from . import add, delete, get, models, prepare, change
|
from . import add, change, delete, get, models, prepare
|
||||||
|
|
||||||
__all__ = ["add", "delete", "get", "models", "prepare", "change"]
|
__all__ = ["add", "delete", "get", "models", "prepare", "change"]
|
||||||
|
@ -5,23 +5,9 @@ from sqlalchemy.future import Engine
|
|||||||
from . import models
|
from . import models
|
||||||
|
|
||||||
|
|
||||||
def add_account(
|
def add_account(engine: Engine, account: models.Account) -> bool:
|
||||||
engine: Engine,
|
|
||||||
user_id: int,
|
|
||||||
name: str,
|
|
||||||
salt: bytes,
|
|
||||||
enc_login: bytes,
|
|
||||||
enc_password: bytes,
|
|
||||||
) -> bool:
|
|
||||||
"""Adds account to the database. Returns true on success,
|
"""Adds account to the database. Returns true on success,
|
||||||
false otherwise"""
|
false otherwise"""
|
||||||
account = models.Account(
|
|
||||||
user_id=user_id,
|
|
||||||
name=name,
|
|
||||||
salt=salt,
|
|
||||||
enc_login=enc_login,
|
|
||||||
enc_password=enc_password,
|
|
||||||
)
|
|
||||||
try:
|
try:
|
||||||
with sqlmodel.Session(engine) as session:
|
with sqlmodel.Session(engine) as session:
|
||||||
session.add(account)
|
session.add(account)
|
||||||
@ -32,19 +18,9 @@ def add_account(
|
|||||||
return True
|
return True
|
||||||
|
|
||||||
|
|
||||||
def add_master_pass(
|
def add_master_pass(engine: Engine, master_pass: models.MasterPass) -> bool:
|
||||||
engine: Engine,
|
|
||||||
user_id: int,
|
|
||||||
salt: bytes,
|
|
||||||
password_hash: bytes,
|
|
||||||
) -> bool:
|
|
||||||
"""Adds master password the database. Returns true on success,
|
"""Adds master password the database. Returns true on success,
|
||||||
false otherwise"""
|
false otherwise"""
|
||||||
master_pass = models.MasterPass(
|
|
||||||
user_id=user_id,
|
|
||||||
salt=salt,
|
|
||||||
password_hash=password_hash,
|
|
||||||
)
|
|
||||||
try:
|
try:
|
||||||
with sqlmodel.Session(engine) as session:
|
with sqlmodel.Session(engine) as session:
|
||||||
session.add(master_pass)
|
session.add(master_pass)
|
||||||
|
@ -5,13 +5,17 @@ from . import models
|
|||||||
|
|
||||||
|
|
||||||
def change_master_pass(
|
def change_master_pass(
|
||||||
engine: Engine, user_id: int, salt: bytes, password: bytes
|
engine: Engine,
|
||||||
|
master_password: models.MasterPass,
|
||||||
) -> None:
|
) -> None:
|
||||||
"""Changes master password and salt in the database"""
|
"""Changes master password and salt in the database"""
|
||||||
statement = (
|
statement = (
|
||||||
sqlmodel.update(models.MasterPass)
|
sqlmodel.update(models.MasterPass)
|
||||||
.where(models.MasterPass.user_id == user_id)
|
.where(models.MasterPass.user_id == master_password.user_id)
|
||||||
.values(salt=salt, password_hash=password)
|
.values(
|
||||||
|
salt=master_password.salt,
|
||||||
|
password_hash=master_password.password_hash,
|
||||||
|
)
|
||||||
)
|
)
|
||||||
with sqlmodel.Session(engine) as session:
|
with sqlmodel.Session(engine) as session:
|
||||||
session.exec(statement)
|
session.exec(statement)
|
||||||
|
@ -1,5 +1,3 @@
|
|||||||
from typing import Iterator
|
|
||||||
|
|
||||||
import sqlmodel
|
import sqlmodel
|
||||||
from sqlalchemy.future import Engine
|
from sqlalchemy.future import Engine
|
||||||
|
|
||||||
@ -9,17 +7,14 @@ from . import models
|
|||||||
def get_master_pass(
|
def get_master_pass(
|
||||||
engine: Engine,
|
engine: Engine,
|
||||||
user_id: int,
|
user_id: int,
|
||||||
) -> tuple[bytes, bytes] | None:
|
) -> models.MasterPass | None:
|
||||||
"""Gets master pass. Returns tuple of salt and password
|
"""Gets master password of a user"""
|
||||||
or None if it wasn't found"""
|
|
||||||
statement = sqlmodel.select(models.MasterPass).where(
|
statement = sqlmodel.select(models.MasterPass).where(
|
||||||
models.MasterPass.user_id == user_id,
|
models.MasterPass.user_id == user_id,
|
||||||
)
|
)
|
||||||
with sqlmodel.Session(engine) as session:
|
with sqlmodel.Session(engine) as session:
|
||||||
result = session.exec(statement).first()
|
result = session.exec(statement).first()
|
||||||
if result is None:
|
return result
|
||||||
return
|
|
||||||
return (result.salt, result.password_hash)
|
|
||||||
|
|
||||||
|
|
||||||
def get_accounts(
|
def get_accounts(
|
||||||
@ -28,7 +23,7 @@ def get_accounts(
|
|||||||
*,
|
*,
|
||||||
to_sort: bool = False,
|
to_sort: bool = False,
|
||||||
) -> list[str]:
|
) -> list[str]:
|
||||||
"""Gets list of account names"""
|
"""Gets a list of account names of a user"""
|
||||||
statement = sqlmodel.select(models.Account.name).where(
|
statement = sqlmodel.select(models.Account.name).where(
|
||||||
models.Account.user_id == user_id,
|
models.Account.user_id == user_id,
|
||||||
)
|
)
|
||||||
@ -39,11 +34,8 @@ def get_accounts(
|
|||||||
return result
|
return result
|
||||||
|
|
||||||
|
|
||||||
def get_all_accounts(
|
def get_all_accounts(engine: Engine, user_id: int) -> list[models.Account]:
|
||||||
engine: Engine, user_id: int
|
"""Returns a list of accounts of a user"""
|
||||||
) -> Iterator[tuple[str, bytes, bytes, bytes]]:
|
|
||||||
"""Returns an iterator of tuples, where values represent account's
|
|
||||||
name, salt, encrypted login and encrypted password"""
|
|
||||||
statement = (
|
statement = (
|
||||||
sqlmodel.select(models.Account)
|
sqlmodel.select(models.Account)
|
||||||
.where(
|
.where(
|
||||||
@ -53,28 +45,19 @@ def get_all_accounts(
|
|||||||
)
|
)
|
||||||
with sqlmodel.Session(engine) as session:
|
with sqlmodel.Session(engine) as session:
|
||||||
result = session.exec(statement).fetchall()
|
result = session.exec(statement).fetchall()
|
||||||
yield from (
|
return result
|
||||||
(
|
|
||||||
account.name,
|
|
||||||
account.salt,
|
|
||||||
account.enc_login,
|
|
||||||
account.enc_password,
|
|
||||||
)
|
|
||||||
for account in result
|
|
||||||
)
|
|
||||||
|
|
||||||
|
|
||||||
def get_account_info(
|
def get_account_info(
|
||||||
engine: Engine, user_id: int, name: str
|
engine: Engine,
|
||||||
) -> tuple[bytes, bytes, bytes]:
|
user_id: int,
|
||||||
"""Gets account info. Returns tuple of salt, login and password
|
name: str,
|
||||||
or None if it wasn't found"""
|
) -> models.Account:
|
||||||
|
"""Gets account info"""
|
||||||
statement = sqlmodel.select(models.Account).where(
|
statement = sqlmodel.select(models.Account).where(
|
||||||
models.Account.user_id == user_id,
|
models.Account.user_id == user_id,
|
||||||
models.Account.name == name,
|
models.Account.name == name,
|
||||||
)
|
)
|
||||||
with sqlmodel.Session(engine) as session:
|
with sqlmodel.Session(engine) as session:
|
||||||
result = session.exec(statement).first()
|
result = session.exec(statement).first()
|
||||||
if result is None:
|
return result
|
||||||
return
|
|
||||||
return (result.salt, result.enc_login, result.enc_password)
|
|
||||||
|
@ -3,6 +3,8 @@ import os
|
|||||||
from cryptography.exceptions import InvalidKey
|
from cryptography.exceptions import InvalidKey
|
||||||
from cryptography.hazmat.primitives.kdf.scrypt import Scrypt
|
from cryptography.hazmat.primitives.kdf.scrypt import Scrypt
|
||||||
|
|
||||||
|
from ..database.models import MasterPass
|
||||||
|
|
||||||
MEMORY_USAGE = 2**14
|
MEMORY_USAGE = 2**14
|
||||||
|
|
||||||
|
|
||||||
@ -17,22 +19,23 @@ def _get_kdf(salt: bytes) -> Scrypt:
|
|||||||
return kdf
|
return kdf
|
||||||
|
|
||||||
|
|
||||||
def encrypt_master_pass(password: str) -> tuple[bytes, bytes]:
|
def encrypt_master_pass(user_id: int, password: str) -> MasterPass:
|
||||||
"""Hashes master password and return tuple of hashed password and salt"""
|
"""Hashes master password and returns MasterPass object"""
|
||||||
salt = os.urandom(64)
|
salt = os.urandom(64)
|
||||||
kdf = _get_kdf(salt)
|
kdf = _get_kdf(salt)
|
||||||
return kdf.derive(password.encode("utf-8")), salt
|
password_hash = kdf.derive(password.encode("utf-8"))
|
||||||
|
return MasterPass(
|
||||||
|
user_id=user_id,
|
||||||
|
password_hash=password_hash,
|
||||||
|
salt=salt,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
def check_master_pass(
|
def check_master_pass(password: str, master_password: MasterPass) -> bool:
|
||||||
password: str,
|
|
||||||
enc_password: bytes,
|
|
||||||
salt: bytes,
|
|
||||||
) -> bool:
|
|
||||||
"""Checks if the master password is correct"""
|
"""Checks if the master password is correct"""
|
||||||
kdf = _get_kdf(salt)
|
kdf = _get_kdf(master_password.salt)
|
||||||
try:
|
try:
|
||||||
kdf.verify(password.encode("utf-8"), enc_password)
|
kdf.verify(password.encode("utf-8"), master_password.password_hash)
|
||||||
except InvalidKey:
|
except InvalidKey:
|
||||||
return False
|
return False
|
||||||
else:
|
else:
|
||||||
|
@ -1,14 +1,18 @@
|
|||||||
import base64
|
import base64
|
||||||
import os
|
import os
|
||||||
from typing import Iterator
|
from typing import Iterable, Iterator
|
||||||
|
|
||||||
from cryptography.fernet import Fernet
|
from cryptography.fernet import Fernet
|
||||||
from cryptography.hazmat.backends import default_backend
|
from cryptography.hazmat.backends import default_backend
|
||||||
from cryptography.hazmat.primitives import hashes
|
from cryptography.hazmat.primitives import hashes
|
||||||
from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
|
from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
|
||||||
|
|
||||||
|
from ..classes import DecryptedAccount
|
||||||
|
from ..database.models import Account
|
||||||
|
|
||||||
|
|
||||||
def _generate_key(salt: bytes, master_pass: bytes) -> bytes:
|
def _generate_key(salt: bytes, master_pass: bytes) -> bytes:
|
||||||
|
"""Generates key for fernet encryption"""
|
||||||
kdf = PBKDF2HMAC(
|
kdf = PBKDF2HMAC(
|
||||||
algorithm=hashes.SHA256(),
|
algorithm=hashes.SHA256(),
|
||||||
length=32,
|
length=32,
|
||||||
@ -21,49 +25,54 @@ def _generate_key(salt: bytes, master_pass: bytes) -> bytes:
|
|||||||
|
|
||||||
|
|
||||||
def encrypt(
|
def encrypt(
|
||||||
login: str,
|
account: DecryptedAccount,
|
||||||
passwd: str,
|
|
||||||
master_pass: str,
|
master_pass: str,
|
||||||
) -> tuple[bytes, bytes, bytes]:
|
) -> Account:
|
||||||
"""Encrypts login and password of a user using their master
|
"""Encrypts account using master password and returns Account object"""
|
||||||
password as a key.
|
|
||||||
Returns a tuple of encrypted login, password and salt"""
|
|
||||||
salt = os.urandom(64)
|
salt = os.urandom(64)
|
||||||
key = _generate_key(salt, master_pass.encode("utf-8"))
|
key = _generate_key(salt, master_pass.encode("utf-8"))
|
||||||
f = Fernet(key)
|
f = Fernet(key)
|
||||||
enc_login = base64.urlsafe_b64decode(f.encrypt(login.encode("utf-8")))
|
enc_login = base64.urlsafe_b64decode(
|
||||||
enc_password = base64.urlsafe_b64decode(f.encrypt(passwd.encode("utf-8")))
|
f.encrypt(account.login.encode("utf-8")),
|
||||||
return (enc_login, enc_password, salt)
|
)
|
||||||
|
enc_password = base64.urlsafe_b64decode(
|
||||||
|
f.encrypt(account.password.encode("utf-8")),
|
||||||
|
)
|
||||||
|
return Account(
|
||||||
|
user_id=account.user_id,
|
||||||
|
name=account.name,
|
||||||
|
salt=salt,
|
||||||
|
enc_login=enc_login,
|
||||||
|
enc_password=enc_password,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
def decrypt(
|
def decrypt(
|
||||||
enc_login: bytes,
|
account: Account,
|
||||||
enc_pass: bytes,
|
|
||||||
master_pass: str,
|
master_pass: str,
|
||||||
salt: bytes,
|
) -> DecryptedAccount:
|
||||||
) -> tuple[str, str]:
|
"""Decrypts account using master password and returns
|
||||||
"""Decrypts login and password using their
|
DecryptedAccount object"""
|
||||||
master password as a key.
|
key = _generate_key(account.salt, master_pass.encode("utf-8"))
|
||||||
Returns a tuple of decrypted login and password"""
|
|
||||||
key = _generate_key(salt, master_pass.encode("utf-8"))
|
|
||||||
f = Fernet(key)
|
f = Fernet(key)
|
||||||
login = f.decrypt(base64.urlsafe_b64encode(enc_login)).decode("utf-8")
|
login = f.decrypt(
|
||||||
password = f.decrypt(base64.urlsafe_b64encode(enc_pass)).decode("utf-8")
|
base64.urlsafe_b64encode(account.enc_login),
|
||||||
return (login, password)
|
).decode("utf-8")
|
||||||
|
password = f.decrypt(
|
||||||
|
base64.urlsafe_b64encode(account.enc_password),
|
||||||
|
).decode("utf-8")
|
||||||
|
return DecryptedAccount(
|
||||||
|
user_id=account.user_id,
|
||||||
|
name=account.name,
|
||||||
|
login=login,
|
||||||
|
password=password,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
def decrypt_multiple(
|
def decrypt_multiple(
|
||||||
accounts: Iterator[tuple[str, bytes, bytes, bytes]], master_pass: str
|
accounts: Iterable[Account], master_pass: str
|
||||||
) -> Iterator[tuple[str, str, str]]:
|
) -> Iterator[DecryptedAccount]:
|
||||||
"""Gets an iterator of tuples, where values represent account's name, salt,
|
"""Decrypts an iterable of accounts using master_pass and
|
||||||
encrypted login and encrypted password.
|
returns an Iterator of decrypted accounts"""
|
||||||
Return an iterator of names, logins and passwords as a tuple"""
|
|
||||||
for account in accounts:
|
for account in accounts:
|
||||||
name, salt, enc_login, enc_passwd = account
|
yield decrypt(account, master_pass)
|
||||||
login, passwd = decrypt(
|
|
||||||
enc_login,
|
|
||||||
enc_passwd,
|
|
||||||
master_pass,
|
|
||||||
salt,
|
|
||||||
)
|
|
||||||
yield (name, login, passwd)
|
|
||||||
|
Reference in New Issue
Block a user