More code cleanup with a few bug fixes

This commit is contained in:
2022-11-30 19:41:57 +03:00
parent 0d3965d5d2
commit 6cd8091fde
9 changed files with 165 additions and 52 deletions

View File

@ -15,16 +15,20 @@ def create_bot(token: str, engine: Engine) -> telebot.TeleBot:
commands=["set_master_pass"],
)
bot.register_message_handler(
functools.partial(handlers.get_account, bot, engine), commands=["get_account"]
functools.partial(handlers.get_account, bot, engine),
commands=["get_account"],
)
bot.register_message_handler(
functools.partial(handlers.get_accounts, bot, engine), commands=["get_accounts"]
functools.partial(handlers.get_accounts, bot, engine),
commands=["get_accounts"],
)
bot.register_message_handler(
functools.partial(handlers.add_account, bot, engine), commands=["add_account"]
functools.partial(handlers.add_account, bot, engine),
commands=["add_account"],
)
bot.register_message_handler(
functools.partial(handlers.delete_all, bot, engine), commands=["delete_all"]
functools.partial(handlers.delete_all, bot, engine),
commands=["delete_all"],
)
bot.register_message_handler(
functools.partial(handlers.reset_master_pass, bot, engine),
@ -35,7 +39,8 @@ def create_bot(token: str, engine: Engine) -> telebot.TeleBot:
commands=["delete_account"],
)
bot.register_message_handler(
functools.partial(handlers.help, bot), commands=["help", "start"]
functools.partial(handlers.help_command, bot),
commands=["help", "start"],
)
bot.register_message_handler(
functools.partial(handlers.cancel, bot), commands=["cancel"]
@ -44,9 +49,11 @@ def create_bot(token: str, engine: Engine) -> telebot.TeleBot:
functools.partial(handlers.export, bot, engine), commands=["export"]
)
bot.register_message_handler(
functools.partial(handlers.import_accounts, bot, engine), commands=["import"]
functools.partial(handlers.import_accounts, bot, engine),
commands=["import"],
)
bot.register_message_handler(
functools.partial(handlers.gen_password, bot), commands=["gen_password"]
functools.partial(handlers.gen_password, bot),
commands=["gen_password"],
)
return bot

View File

@ -58,7 +58,8 @@ def delete_all(bot: telebot.TeleBot, engine: Engine, mes: Message) -> None:
_base_handler(bot, mes)
bot_mes = bot.send_message(
mes.chat.id,
"Вы действительно хотите удалить все ваши аккаунты? Это действие нельзя отменить. "
"Вы действительно хотите удалить все ваши аккаунты? Это действие "
"нельзя отменить. "
"Отправьте YES для подтверждения",
)
bot.register_next_step_handler(
@ -76,13 +77,25 @@ def _delete_all(
database.delete.delete_master_pass(engine, mes.from_user.id)
_send_tmp_message(bot, mes.chat.id, "Всё успешно удалено", timeout=10)
else:
_send_tmp_message(bot, mes.chat.id, "Вы отправили не YES, ничего не удалено")
_send_tmp_message(
bot,
mes.chat.id,
"Вы отправили не YES, ничего не удалено",
)
def set_master_password(bot: telebot.TeleBot, engine: Engine, mes: Message) -> None:
def set_master_password(
bot: telebot.TeleBot,
engine: Engine,
mes: Message,
) -> None:
_base_handler(bot, mes, None)
if database.get.get_master_pass(engine, mes.from_user.id) is not 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(
mes, functools.partial(_set_master_pass2, bot, engine, bot_mes)
@ -97,14 +110,25 @@ def _set_master_pass2(
if text == "/cancel":
return _send_tmp_message(bot, mes.chat.id, "Успешная отмена")
hash_pass, master_salt = cryptography.master_pass.encrypt_master_pass(text)
database.add.add_master_pass(engine, mes.from_user.id, master_salt, hash_pass)
password_hash, master_salt = cryptography.master_pass.encrypt_master_pass(
text,
)
database.add.add_master_pass(
engine,
mes.from_user.id,
master_salt,
password_hash,
)
_send_tmp_message(bot, mes.chat.id, "Успех")
del mes, text
gc.collect()
def reset_master_pass(bot: telebot.TeleBot, engine: Engine, mes: Message) -> None:
def reset_master_pass(
bot: telebot.TeleBot,
engine: Engine,
mes: Message,
) -> None:
_base_handler(bot, mes)
if database.get.get_master_pass(engine, mes.from_user.id) is None:
return _send_tmp_message(bot, mes.chat.id, "Мастер пароль не задан")
@ -139,7 +163,10 @@ def _reset_master_pass2(
def add_account(bot: telebot.TeleBot, engine: Engine, mes: Message) -> None:
_base_handler(bot, mes)
master_password_from_db = database.get.get_master_pass(engine, mes.from_user.id)
master_password_from_db = database.get.get_master_pass(
engine,
mes.from_user.id,
)
if master_password_from_db is None:
return _send_tmp_message(bot, mes.chat.id, "Нет мастер пароля")
@ -159,7 +186,11 @@ def _add_account2(
return _send_tmp_message(bot, mes.chat.id, "Успешная отмена")
if not check_account_name(text):
return _send_tmp_message(bot, mes.chat.id, "Не корректное название аккаунта")
return _send_tmp_message(
bot,
mes.chat.id,
"Не корректное название аккаунта",
)
if text in database.get.get_accounts(engine, mes.from_user.id):
return _send_tmp_message(
bot, mes.chat.id, "Аккаунт с таким именем уже существует"
@ -233,12 +264,18 @@ def _add_account5(
salt, hash_ = database.get.get_master_pass(engine, mes.from_user.id)
if not cryptography.master_pass.check_master_pass(text, hash_, salt):
return _send_tmp_message(bot, mes.chat.id, "Не подходит главный пароль")
return _send_tmp_message(
bot,
mes.chat.id,
"Не подходит главный пароль",
)
name, login, passwd = data["name"], data["login"], data["passwd"]
enc_login, enc_pass, salt = cryptography.other_accounts.encrypt_account_info(
login, passwd, text
enc_login, enc_pass, salt = cryptography.other_accounts.encrypt(
login,
passwd,
text,
)
result = database.add.add_account(
@ -246,7 +283,9 @@ def _add_account5(
)
_send_tmp_message(
bot, mes.chat.id, "Успех" if result else "Произошла не предвиденная ошибка"
bot,
mes.chat.id,
"Успех" if result else "Произошла не предвиденная ошибка",
)
del data, name, login, passwd, enc_login
@ -285,29 +324,43 @@ def _get_account2(
def _get_account3(
bot: telebot.TeleBot, engine: Engine, prev_mes: Message, name: str, mes: Message
bot: telebot.TeleBot,
engine: Engine,
prev_mes: Message,
name: str,
mes: Message,
) -> None:
_base_handler(bot, mes, prev_mes)
text = mes.text.strip()
if text == "/cancel":
return _send_tmp_message(bot, mes.chat.id, "Успешная отмена")
master_salt, hash_pass = database.get.get_master_pass(engine, mes.from_user.id)
master_salt, hash_pass = database.get.get_master_pass(
engine,
mes.from_user.id,
)
if not cryptography.master_pass.check_master_pass(text, hash_pass, master_salt):
if not cryptography.master_pass.check_master_pass(
text,
hash_pass,
master_salt,
):
return _send_tmp_message(bot, mes.chat.id, "Не подходит мастер пароль")
salt, enc_login, enc_pass = database.get.get_account_info(
engine, mes.from_user.id, name
)
login, passwd = cryptography.other_accounts.decrypt_account_info(
enc_login, enc_pass, text, salt
login, passwd = cryptography.other_accounts.decrypt(
enc_login,
enc_pass,
text,
salt,
)
_send_tmp_message(
bot,
mes.chat.id,
f"Логин:\n`{login}`\nПароль:\n`{passwd}`\nНажмите на логин или пароль, "
"чтобы скопировать",
f"Логин:\n`{login}`\nПароль:\n`{passwd}`\nНажмите на логин "
"или пароль, чтобы скопировать",
30,
)
@ -346,7 +399,7 @@ def _delete_account2(
_send_tmp_message(bot, mes.chat.id, "Аккаунт удалён")
def help(bot: telebot.TeleBot, mes: Message) -> None:
def help_command(bot: telebot.TeleBot, mes: Message) -> None:
message = """Команды:
/set_master_pass - установить мастер пароль
/add_account - создать аккаунт
@ -371,7 +424,10 @@ def cancel(bot: telebot.TeleBot, mes: Message) -> None:
def export(bot: telebot.TeleBot, engine: Engine, mes: Message) -> None:
_base_handler(bot, mes)
master_password_from_db = database.get.get_master_pass(engine, mes.from_user.id)
master_password_from_db = database.get.get_master_pass(
engine,
mes.from_user.id,
)
if master_password_from_db is None:
return _send_tmp_message(bot, mes.chat.id, "Нет мастер пароля")
@ -394,8 +450,15 @@ def _export2(
if text == "/cancel":
return _send_tmp_message(bot, mes.chat.id, "Успешная отмена")
master_salt, hash_pass = database.get.get_master_pass(engine, mes.from_user.id)
if not cryptography.master_pass.check_master_pass(text, hash_pass, master_salt):
master_salt, hash_pass = database.get.get_master_pass(
engine,
mes.from_user.id,
)
if not cryptography.master_pass.check_master_pass(
text,
hash_pass,
master_salt,
):
return _send_tmp_message(bot, mes.chat.id, "Не подходит мастер пароль")
accounts = database.get.get_all_accounts(engine, mes.from_user.id)
@ -409,9 +472,16 @@ def _export2(
bot.delete_message(bot_mes.chat.id, bot_mes.id)
def import_accounts(bot: telebot.TeleBot, engine: Engine, mes: Message) -> None:
def import_accounts(
bot: telebot.TeleBot,
engine: Engine,
mes: Message,
) -> None:
_base_handler(bot, mes)
master_password_from_db = database.get.get_master_pass(engine, mes.from_user.id)
master_password_from_db = database.get.get_master_pass(
engine,
mes.from_user.id,
)
if master_password_from_db is None:
return _send_tmp_message(bot, mes.chat.id, "Нет мастер пароля")
@ -433,7 +503,11 @@ def _import2(
return _send_tmp_message(bot, mes.chat.id, "Успешная отмена")
if mes.document is None:
return _send_tmp_message(bot, mes.chat.id, "Вы должны отправить документ")
return _send_tmp_message(
bot,
mes.chat.id,
"Вы должны отправить документ",
)
if mes.document.file_size > 102_400: # If file size is bigger that 100 MB
return _send_tmp_message(bot, mes.chat.id, "Файл слишком большой")
@ -442,7 +516,11 @@ def _import2(
try:
accounts = json_to_accounts(downloaded_file.decode("utf-8"))
except Exception:
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(
@ -462,18 +540,26 @@ def _import3(
if text == "/cancel":
return _send_tmp_message(bot, mes.chat.id, "Успешная отмена")
master_salt, hash_pass = database.get.get_master_pass(engine, mes.from_user.id)
if not cryptography.master_pass.check_master_pass(text, hash_pass, master_salt):
master_salt, hash_pass = database.get.get_master_pass(
engine,
mes.from_user.id,
)
if not cryptography.master_pass.check_master_pass(
text,
hash_pass,
master_salt,
):
return _send_tmp_message(bot, mes.chat.id, "Не подходит мастер пароль")
# List of names of accounts, which failed to be added to the database or failed tests
# List of names of accounts, which failed to be added to the database
# or failed the tests
failed: list[str] = []
for account in accounts:
name, login, passwd = account
if not check_account(name, login, passwd):
failed.append(name)
continue
enc_login, enc_passwd, salt = cryptography.other_accounts.encrypt_account_info(
enc_login, enc_passwd, salt = cryptography.other_accounts.encrypt(
login, passwd, text
)
result = database.add.add_account(