diff --git a/README.md b/README.md index 14f3b85..57c0bf2 100644 --- a/README.md +++ b/README.md @@ -22,6 +22,7 @@ - /cancel - отмена текущего действия - /help - помощь - /export - получить пароли в json формате +- /import - импортировать пароли из json в файле в таком же формате, как из /export ### Настройка diff --git a/src/bot/__init__.py b/src/bot/__init__.py index 12df942..c8f2914 100644 --- a/src/bot/__init__.py +++ b/src/bot/__init__.py @@ -43,4 +43,7 @@ def create_bot(token: str, engine: mariadb.Connection) -> telebot.TeleBot: bot.register_message_handler( functools.partial(handlers.export, bot, engine), commands=["export"] ) + bot.register_message_handler( + functools.partial(handlers.import_accounts, bot, engine), commands=["import"] + ) return bot diff --git a/src/bot/handlers.py b/src/bot/handlers.py index 903e199..0c1c359 100644 --- a/src/bot/handlers.py +++ b/src/bot/handlers.py @@ -6,7 +6,13 @@ import telebot from sqlalchemy.future import Engine from .. import cryptography, database -from .utils import accounts_to_json, base_handler, get_all_accounts, send_tmp_message +from .utils import ( + accounts_to_json, + base_handler, + get_all_accounts, + json_to_accounts, + send_tmp_message, +) Message = telebot.types.Message @@ -315,7 +321,8 @@ def help(bot: telebot.TeleBot, mes: telebot.types.Message) -> None: /reset_master_pass - удалить все аккаунты и изменить мастер пароль /cancel - отмена текущего действия /help - помощь -/export - получить пароли в json формате""" +/export - получить пароли в json формате +/import - импортировать пароли из json в файле в таком же формате, как из /export""" bot.send_message(mes.chat.id, message) @@ -360,3 +367,79 @@ def _export2( gc.collect() time.sleep(30) bot.delete_message(bot_mes.chat.id, bot_mes.id) + + +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) + + if master_password_from_db is None: + return send_tmp_message(bot, mes.chat.id, "Нет мастер пароля") + + bot_mes = bot.send_message(mes.chat.id, "Отправьте json файл") + + bot.register_next_step_handler( + mes, functools.partial(_import2, bot, engine, bot_mes) + ) + + +def _import2( + bot: telebot.TeleBot, engine: Engine, prev_mes: Message, mes: Message +) -> None: + base_handler(bot, mes, prev_mes) + if mes.text is not None: + text = mes.text.strip() + if text == "/cancel": + return send_tmp_message(bot, mes.chat.id, "Успешная отмена") + + if mes.document is None: + send_tmp_message(bot, mes.chat.id, "Вы должны отправить документ") + + file_info = bot.get_file(mes.document.file_id) + downloaded_file = bot.download_file(file_info.file_path) + try: + accounts = json_to_accounts(downloaded_file.decode("utf-8")) + except Exception: + send_tmp_message(bot, mes.chat.id, "Ошибка во время работы с файлом") + + bot_mes = bot.send_message(mes.chat.id, "Отправьте мастер пароль") + bot.register_next_step_handler( + mes, functools.partial(_import3, bot, engine, bot_mes, accounts) + ) + + +def _import3( + bot: telebot.TeleBot, + engine: Engine, + prev_mes: Message, + accounts: list[tuple[str, str, 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) + if cryptography.master_pass.encrypt_master_pass(text, master_salt) != hash_pass: + return send_tmp_message(bot, mes.chat.id, "Не подходит мастер пароль") + + failed: list[str] = [] + for account in accounts: + name, login, passwd = account + enc_login, enc_passwd, salt = cryptography.other_accounts.encrypt_account_info( + login, passwd, text.encode("utf-8") + ) + result = database.add.add_account( + engine, mes.from_user.id, name, salt, enc_login, enc_passwd + ) + if not result: + failed.append(name) + + if failed: + mes_text = "Не удалось добавить:\n" + "\n".join(failed) + else: + mes_text = "Успех" + send_tmp_message(bot, mes.chat.id, mes_text, 10) + del text, mes, accounts + gc.collect()