Now commands ask for params in separate messages

This commit is contained in:
StNicolay 2022-10-30 00:31:51 +03:00
parent fae04547c8
commit dec7a9b7c9
2 changed files with 217 additions and 119 deletions

View File

@ -22,7 +22,7 @@ def create_bot(token: str, engine: mariadb.Connection) -> telebot.TeleBot:
functools.partial(handlers.get_accounts, bot, engine), commands=["get_accounts"] functools.partial(handlers.get_accounts, bot, engine), commands=["get_accounts"]
) )
bot.register_message_handler( bot.register_message_handler(
functools.partial(handlers.add_record, bot, engine), commands=["add_account"] functools.partial(handlers.add_account, bot, engine), commands=["add_account"]
) )
bot.register_message_handler( bot.register_message_handler(
functools.partial(handlers.delete_all, bot, engine), commands=["delete_all"] functools.partial(handlers.delete_all, bot, engine), commands=["delete_all"]

View File

@ -1,5 +1,5 @@
import functools
import gc import gc
import shlex
import time import time
import telebot import telebot
@ -7,6 +7,8 @@ from sqlalchemy.future import Engine
from .. import cryptography, database from .. import cryptography, database
Message = telebot.types.Message
def _send_tmp_message( def _send_tmp_message(
bot: telebot.TeleBot, chat_id: telebot.types.Message, text: str, timeout: int = 5 bot: telebot.TeleBot, chat_id: telebot.types.Message, text: str, timeout: int = 5
@ -16,41 +18,6 @@ def _send_tmp_message(
bot.delete_message(chat_id, bot_mes.id) bot.delete_message(chat_id, bot_mes.id)
def add_record(
bot: telebot.TeleBot, engine: Engine, mes: telebot.types.Message
) -> None:
data = shlex.split(mes.text)
bot.delete_message(mes.chat.id, mes.id)
if len(data) != 5:
return _send_tmp_message(bot, mes.chat.id, "Неправильное кол-во аргументов")
master_password = data[4]
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, "Нет мастер пароля")
master_salt, hash_pass = master_password_from_db
if (
cryptography.master_pass.encrypt_master_pass(master_password, master_salt)
!= hash_pass
):
return _send_tmp_message(bot, mes.chat.id, "Не подходит главный пароль")
enc_login, enc_pass, salt = cryptography.other_accounts.encrypt_account_info(
data[2], data[3], master_password.encode("utf-8")
)
result = database.add.add_account(
engine, mes.from_user.id, data[1], salt, enc_login, enc_pass
)
_send_tmp_message(
bot,
mes.chat.id,
"Успех" if result else "Ошибка, вероятно аккаунт с таким же именем существует",
)
del data, master_password, mes
gc.collect()
def get_accounts( def get_accounts(
bot: telebot.TeleBot, engine: Engine, mes: telebot.types.Message bot: telebot.TeleBot, engine: Engine, mes: telebot.types.Message
) -> None: ) -> None:
@ -61,65 +28,10 @@ def get_accounts(
bot, bot,
mes.chat.id, mes.chat.id,
"Ваши аккаунты:\n" + "\n".join(accounts) if accounts else "У вас нет аккаунтов", "Ваши аккаунты:\n" + "\n".join(accounts) if accounts else "У вас нет аккаунтов",
timeout=30,
) )
def set_master_password(
bot: telebot.TeleBot, engine: Engine, mes: telebot.types.Message
) -> None:
data = shlex.split(mes.text)
bot.delete_message(mes.chat.id, mes.id)
if database.get.get_master_pass(engine, mes.from_user.id) is not None:
return _send_tmp_message(bot, mes.chat.id, "У вас уже установлен мастер пароль")
if len(data) != 2:
return _send_tmp_message(bot, mes.chat.id, "Неправильное количество аргументов")
pass_, salt = cryptography.master_pass.encrypt_master_pass(data[1])
result = database.add.add_master_pass(engine, mes.from_user.id, salt, pass_)
_send_tmp_message(
bot,
mes.chat.id,
"Успех" if result else "Ошибка, вероятно аккаунт с таким же именем существует",
)
del data, mes
gc.collect()
def get_account(
bot: telebot.TeleBot, engine: Engine, mes: telebot.types.Message
) -> None:
data = shlex.split(mes.text)
bot.delete_message(mes.chat.id, mes.id)
if len(data) != 3:
return _send_tmp_message(bot, mes.chat.id, "Неправильное количество аргументов")
master_pass = database.get.get_master_pass(engine, mes.from_user.id)
if master_pass is None:
return _send_tmp_message(bot, mes.chat.id, "Нет мастер пароля")
master_salt, hash_pass = master_pass
if data[1] not in database.get.get_accounts(engine, mes.from_user.id):
return _send_tmp_message(bot, mes.chat.id, "Нет такого аккаунта")
if cryptography.master_pass.encrypt_master_pass(data[2], master_salt) != hash_pass:
return _send_tmp_message(bot, mes.chat.id, "Не подходит главный пароль")
salt, enc_login, enc_pass = database.get.get_account_info(
engine, mes.from_user.id, data[1]
)
login, passwd = cryptography.other_accounts.decrypt_account_info(
enc_login, enc_pass, data[2].encode("utf-8"), salt
)
_send_tmp_message(bot, mes.chat.id, f"Логин:\n{login}\nПароль:\n{passwd}", 30)
del data, mes, passwd, login
gc.collect()
def delete_all( def delete_all(
bot: telebot.TeleBot, engine: Engine, mes: telebot.types.Message bot: telebot.TeleBot, engine: Engine, mes: telebot.types.Message
) -> None: ) -> None:
@ -128,50 +40,236 @@ def delete_all(
bot.send_message(mes.chat.id, "Все ваши данные удалены из базы данных") bot.send_message(mes.chat.id, "Все ваши данные удалены из базы данных")
def reset_master_pass( def _base(bot: telebot.TeleBot, mes: Message, prev_mes: Message | None = None) -> None:
bot: telebot.TeleBot, engine: Engine, mes: telebot.types.Message
) -> None:
bot.delete_message(mes.chat.id, mes.id) bot.delete_message(mes.chat.id, mes.id)
data = shlex.split(mes.text) if prev_mes is not None:
if len(data) != 2: bot.delete_message(prev_mes.chat.id, prev_mes.id)
return _send_tmp_message(bot, mes.chat.id, "Неправильное количество аргументов")
def set_master_password(bot: telebot.TeleBot, engine: Engine, mes: Message) -> None:
_base(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, "Мастер пароль уже существует")
bot_mes = bot.send_message(mes.chat.id, "Отправьте мастер пароль")
bot.register_next_step_handler(
mes, functools.partial(_set_master_pass2, bot, engine, bot_mes)
)
def _set_master_pass2(
bot: telebot.TeleBot, engine: Engine, prev_mes: Message, mes: Message
) -> None:
_base(bot, mes, prev_mes)
hash_, salt = cryptography.master_pass.encrypt_master_pass(mes.text.strip())
database.add.add_master_pass(engine, mes.from_user.id, salt, hash_)
_send_tmp_message(bot, mes.chat.id, "Успех")
del mes
gc.collect()
def reset_master_pass(bot: telebot.TeleBot, engine: Engine, mes: Message) -> None:
_base(bot, mes)
if database.get.get_master_pass(engine, mes.from_user.id) is None:
return _send_tmp_message(bot, mes.chat.id, "Мастер пароль не задан")
bot_mes = bot.send_message(mes.chat.id, "Отправьте мастер пароль")
bot.register_next_step_handler(
mes, functools.partial(_reset_master_pass2, bot, engine, bot_mes)
)
def _reset_master_pass2(
bot: telebot.TeleBot, engine: Engine, prev_mes: Message, mes: Message
) -> None:
_base(bot, mes, prev_mes)
if database.get.get_master_pass(engine, mes.from_user.id) is None:
return _send_tmp_message(bot, mes.chat.id, "Мастер пароль не задан")
hash_, salt = cryptography.master_pass.encrypt_master_pass(mes.text.strip())
database.delete.purge_accounts(engine, mes.from_user.id)
database.change.change_master_pass(engine, mes.from_user.id, salt, hash_)
_send_tmp_message(
bot, mes.chat.id, "Все ваши аккаунты удалены, а мастер пароль изменён"
)
del mes
gc.collect()
def add_account(bot: telebot.TeleBot, engine: Engine, mes: Message) -> None:
_base(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: if master_password_from_db is None:
return _send_tmp_message(bot, mes.chat.id, "Нет мастер пароля") return _send_tmp_message(bot, mes.chat.id, "Нет мастер пароля")
master_password = data[1] bot_mes = bot.send_message(mes.chat.id, "Отправьте название аккаунта")
enc_pass, salt = cryptography.master_pass.encrypt_master_pass(master_password)
database.delete.purge_accounts(engine, mes.from_user.id) bot.register_next_step_handler(
database.change.change_master_pass(engine, mes.from_user.id, salt, enc_pass) mes, functools.partial(_add_account2, bot, engine, bot_mes)
_send_tmp_message(
bot, mes.chat.id, "Все ваши аккаунты удалены, а мастер пароль изменён"
) )
def delete_account( def _add_account2(
bot: telebot.TeleBot, engine: Engine, mes: telebot.types.Message bot: telebot.TeleBot, engine: Engine, prev_mes: Message, mes: Message
) -> None: ) -> None:
data = shlex.split(mes.text) _base(bot, mes, prev_mes)
bot.delete_message(mes.chat.id, mes.id)
if len(data) != 2: if (
return _send_tmp_message(bot, mes.chat.id, "Неправильное количество аргументов") database.get.get_account_info(engine, mes.from_user.id, mes.text.strip())
is not None
):
_send_tmp_message(bot, mes.chat.id, "Аккаунт с таким именем уже существует")
if data[1] not in database.get.get_accounts(engine, mes.from_user.id): bot_mes = bot.send_message(mes.chat.id, "Отправьте логин")
data = {"name": mes.text.strip()}
bot.register_next_step_handler(
mes, functools.partial(_add_account3, bot, engine, bot_mes, data)
)
def _add_account3(
bot: telebot.TeleBot,
engine: Engine,
prev_mes: Message,
data: dict[str, str],
mes: Message,
) -> None:
_base(bot, mes, prev_mes)
data["login"] = mes.text.strip()
bot_mes = bot.send_message(mes.chat.id, "Отправьте пароль от аккаунта")
bot.register_next_step_handler(
mes, functools.partial(_add_account4, bot, engine, bot_mes, data)
)
def _add_account4(
bot: telebot.TeleBot,
engine: Engine,
prev_mes: Message,
data: dict[str, str],
mes: Message,
) -> None:
_base(bot, mes, prev_mes)
data["passwd"] = mes.text.strip()
bot_mes = bot.send_message(mes.chat.id, "Отправьте мастер пароль")
bot.register_next_step_handler(
mes, functools.partial(_add_account5, bot, engine, bot_mes, data)
)
def _add_account5(
bot: telebot.TeleBot,
engine: Engine,
prev_mes: Message,
data: dict[str, str],
mes: Message,
) -> None:
_base(bot, mes, prev_mes)
text = mes.text.strip()
salt, hash_ = database.get.get_master_pass(engine, mes.from_user.id)
if cryptography.master_pass.encrypt_master_pass(text, salt) != hash_:
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.encode("utf-8")
)
database.add.add_account(engine, mes.from_user.id, name, salt, enc_login, enc_pass)
_send_tmp_message(bot, mes.chat.id, "Успех")
del data, name, login, passwd, enc_login
gc.collect()
def get_account(bot: telebot.TeleBot, engine: Engine, mes: Message) -> None:
_base(bot, mes)
bot_mes = bot.send_message(mes.chat.id, "Отправьте название аккаунта")
bot.register_next_step_handler(
mes, functools.partial(_get_account2, bot, engine, bot_mes)
)
def _get_account2(
bot: telebot.TeleBot, engine: Engine, prev_mes: Message, mes: Message
) -> None:
_base(bot, mes, prev_mes)
text = mes.text.strip()
if text not in database.get.get_accounts(engine, mes.from_user.id):
return _send_tmp_message(bot, mes.chat.id, "Нет такого аккаунта") return _send_tmp_message(bot, mes.chat.id, "Нет такого аккаунта")
database.delete.delete_account(engine, mes.from_user.id, data[1]) bot_mes = bot.send_message(mes.chat.id, "Отправьте мастер пароль")
bot.register_next_step_handler(
mes, functools.partial(_get_account3, bot, engine, bot_mes, text)
)
def _get_account3(
bot: telebot.TeleBot, engine: Engine, prev_mes: Message, name: str, mes: Message
) -> None:
_base(bot, mes, prev_mes)
text = mes.text.strip()
master_pass = database.get.get_master_pass(engine, mes.from_user.id)
if master_pass is None:
return _send_tmp_message(bot, mes.chat.id, "Нет мастер пароля")
master_salt, hash_pass = master_pass
if cryptography.master_pass.encrypt_master_pass(text, master_salt) != hash_pass:
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.encode("utf-8"), salt
)
_send_tmp_message(bot, mes.chat.id, f"Логин:\n{login}\nПароль:\n{passwd}", 30)
del text, mes, passwd, login
gc.collect()
def delete_account(bot: telebot.TeleBot, engine: Engine, mes: Message) -> None:
_base(bot, mes)
bot_mes = bot.send_message(
mes.chat.id, "Отправьте название аккаунта, который вы хотите удалить"
)
bot.register_next_step_handler(
mes, functools.partial(_delete_account2, bot, engine, bot_mes)
)
def _delete_account2(
bot: telebot.TeleBot, engine: Engine, prev_mes: Message, mes: Message
) -> None:
_base(bot, mes, prev_mes)
text = mes.text.strip()
if text not in database.get.get_accounts(engine, mes.from_user.id):
return _send_tmp_message(bot, mes.chat.id, "Нет такого аккаунта")
database.delete.delete_account(engine, mes.from_user.id, text)
_send_tmp_message(bot, mes.chat.id, "Аккаунт удалён") _send_tmp_message(bot, mes.chat.id, "Аккаунт удалён")
def help(bot: telebot.TeleBot, mes: telebot.types.Message) -> None: def help(bot: telebot.TeleBot, mes: telebot.types.Message) -> None:
message = """Команды: message = """Команды:
/set_master_pass {Мастер пароль} - установить мастер пароль /set_master_pass - установить мастер пароль
/add_account {Название} {Логин} {Пароль} {Мастер пароль} - создать аккаунт /add_account - создать аккаунт
/get_accounts - получить список аккаунтов /get_accounts - получить список аккаунтов
/get_account {Название} {Мастер пароль} - получить логин и пароль аккаунта /get_account - получить логин и пароль аккаунта
/delete_account {Название} - удалить аккаунт /delete_account - удалить аккаунт
/delete_all - удалить все аккаунты и мастер пароль /delete_all - удалить все аккаунты и мастер пароль (Спать хочу -> подтверждения нету)
/reset_master_pass {Новый мастер пароль} - удалить все аккаунты и изменить мастер пароль""" /reset_master_pass- удалить все аккаунты и изменить мастер пароль"""
bot.send_message(mes.chat.id, message) bot.send_message(mes.chat.id, message)