Compare commits

...

14 Commits

Author SHA1 Message Date
74844da4ae Added backslash to the forbidden chars 2022-12-30 13:54:03 +03:00
77be64ed4b sorted imports 2022-12-29 18:22:37 +03:00
5bec51beb2 Added keyboard markup to /delete_account 2022-12-29 15:42:27 +03:00
d5d87a8f3b Quality of life changes
Changed default action of /cancel to be handled at message_handler
Added keyboard markup to get_account for easier account selection
2022-12-29 15:38:38 +03:00
d4c50432d7 Added ability to send message, which can be deleted by pressing a button 2022-12-29 15:25:20 +03:00
0026e3321a Made handlers asynchrounos
Switched from TeleBot class to AsyncTeleBot
2022-12-29 13:58:53 +03:00
8858aa09a7 Added conformation before deleting an account. No longer raising an exception if message deletion fails 2022-12-25 21:34:09 +03:00
e165020111 Added blank lines in encryption.other_accounts for better readability 2022-12-25 21:04:45 +03:00
50eb3057d5 Renamed classes.py into decrypted_accounts.py. Removed unnessesary elements from that file 2022-12-25 21:01:18 +03:00
3f744723a9 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
2022-12-25 20:12:19 +03:00
bbc9650357 Changed database.get file. Added fetchall method calls to close session sooner 2022-12-25 17:59:26 +03:00
f299173e56 Changes to database.get_accounts
Function is only fetching account names from db and not whole accounts now
Made sorting in that function optional
2022-12-25 17:35:37 +03:00
5991041b35 Added timeout to the deletion of the message in _send_tmp_message 2022-12-16 06:25:00 +00:00
c2eca49933 Switched to building an image localy instead of using one from docker hub
Added database folder
Created Dockerfile and config for mariadb image
Moved data folder tp database/data/
2022-12-15 15:13:39 +00:00
23 changed files with 1035 additions and 752 deletions

View File

@ -24,5 +24,5 @@
**/secrets.dev.yaml
**/values.dev.yaml
README.md
data/
database/
.flake8

2
.gitignore vendored
View File

@ -161,4 +161,4 @@ cython_debug/
#.idea/
# Database data
data/
database/data/

View File

@ -20,7 +20,9 @@ services:
- password_manager
db:
image: stnicolay/mariadb-aria
build:
context: ./database/
dockerfile: ./Dockerfile
restart: always
environment:
MYSQL_ROOT_PASSWORD: example123!
@ -28,6 +30,6 @@ services:
MYSQL_USER: manager
MYSQL_PASSWORD: passwd123!
volumes:
- ./data:/var/lib/mysql
- ./database/data:/var/lib/mysql
networks:
- password_manager

3
database/Dockerfile Normal file
View File

@ -0,0 +1,3 @@
FROM mariadb
COPY mariadb.cnf /etc/mysql/mariadb.cnf

38
database/mariadb.cnf Normal file
View File

@ -0,0 +1,38 @@
# The MariaDB configuration file
#
# The MariaDB/MySQL tools read configuration files in the following order:
# 0. "/etc/mysql/my.cnf" symlinks to this file, reason why all the rest is read.
# 1. "/etc/mysql/mariadb.cnf" (this file) to set global defaults,
# 2. "/etc/mysql/conf.d/*.cnf" to set global options.
# 3. "/etc/mysql/mariadb.conf.d/*.cnf" to set MariaDB-only options.
# 4. "~/.my.cnf" to set user-specific options.
#
# If the same option is defined multiple times, the last one will apply.
#
# One can use all long options that the program supports.
# Run program with --help to get a list of available options and with
# --print-defaults to see which it would actually understand and use.
#
# If you are new to MariaDB, check out https://mariadb.com/kb/en/basic-mariadb-articles/
#
# This group is read both by the client and the server
# use it for options that affect everything
#
[client-server]
# Port or socket location where to connect
# port = 3306
socket = /run/mysqld/mysqld.sock
# Import all .cnf files from configuration directory
[mariadbd]
skip-host-cache
skip-name-resolve
[mysqld]
skip-innodb
default-storage-engine=Aria
default-tmp-storage-engine=Aria
!includedir /etc/mysql/mariadb.conf.d/
!includedir /etc/mysql/conf.d/

View File

@ -1,2 +1,3 @@
black
flake8
isort

View File

@ -4,3 +4,4 @@ python-dotenv
pyTelegramBotAPI
sqlmodel
pydantic
aiohttp

View File

@ -1,3 +1,4 @@
import asyncio
import os
from dotenv import load_dotenv
@ -6,8 +7,9 @@ from . import (
account_checks,
account_parsing,
bot,
encryption,
database,
decrypted_account,
encryption,
generate_password,
)
@ -15,6 +17,7 @@ __all__ = [
"account_checks",
"account_parsing",
"bot",
"decrypted_account",
"encryption",
"database",
"generate_password",
@ -31,4 +34,4 @@ def main() -> None:
)
database.prepare.prepare(engine)
bot_ = bot.create_bot(os.getenv("TG_TOKEN"), engine)
bot_.infinity_polling()
asyncio.run(bot_.infinity_polling())

View File

@ -1,6 +1,8 @@
import string
FORBIDDEN_CHARS = frozenset("`\n")
from .decrypted_account import DecryptedAccount
FORBIDDEN_CHARS = frozenset("`\n\\")
PUNCTUATION = frozenset(string.punctuation).difference(FORBIDDEN_CHARS)
@ -24,13 +26,13 @@ def check_password(passwd: str) -> bool:
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"""
return all(
(
check_account_name(name),
check_login(login),
check_password(passwd),
check_account_name(account.name),
check_login(account.login),
check_password(account.password),
)
)

View File

@ -1,37 +1,50 @@
import io
from typing import Iterator, Self, Type
from typing import Iterable, Self
import pydantic
from .decrypted_account import DecryptedAccount
class _Account(pydantic.BaseModel):
name: str
login: str
password: str
@classmethod
def from_tuple(cls: Type[Self], tuple_: tuple[str, str, str]) -> Self:
return cls(name=tuple_[0], login=tuple_[1], password=tuple_[2])
def to_usual_account(self, user_id: int) -> DecryptedAccount:
return DecryptedAccount(
user_id=user_id,
name=self.name,
login=self.login,
password=self.password,
)
def as_tuple(self: Self) -> tuple[str, str, str]:
return (self.name, self.login, self.password)
@classmethod
def from_usual_account(cls, account: DecryptedAccount) -> Self:
return cls(
name=account.name,
login=account.login,
password=account.password,
)
class _Accounts(pydantic.BaseModel):
accounts: list[_Account] = pydantic.Field(default_factory=list)
def _accounts_list_to_json(accounts: Iterator[tuple[str, str, str]]) -> str:
accounts = _Accounts(accounts=[_Account.from_tuple(i) for i in accounts])
return accounts.json(ensure_ascii=False)
def _accounts_list_to_json(accounts: Iterable[DecryptedAccount]) -> str:
result = _Accounts(
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_)
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.name = "passwords.json"
return file

View File

@ -1,59 +1,77 @@
import functools
import telebot
from sqlalchemy.future import Engine
from telebot.async_telebot import AsyncTeleBot
from . import handlers
from . import callback_handlers, message_handlers
__all__ = ["handlers"]
__all__ = ["callback_handlers", "message_handlers"]
def create_bot(token: str, engine: Engine) -> telebot.TeleBot:
bot = telebot.TeleBot(token)
def create_bot(token: str, engine: Engine) -> AsyncTeleBot:
bot = AsyncTeleBot(token)
bot.register_message_handler(
functools.partial(handlers.set_master_password, bot, engine),
functools.partial(message_handlers.set_master_password, bot, engine),
commands=["set_master_pass"],
)
bot.register_message_handler(
functools.partial(handlers.get_account, bot, engine),
functools.partial(message_handlers.get_account, bot, engine),
commands=["get_account"],
)
bot.register_message_handler(
functools.partial(handlers.get_accounts, bot, engine),
functools.partial(message_handlers.get_accounts, bot, engine),
commands=["get_accounts"],
)
bot.register_message_handler(
functools.partial(handlers.add_account, bot, engine),
functools.partial(message_handlers.add_account, bot, engine),
commands=["add_account"],
)
bot.register_message_handler(
functools.partial(handlers.delete_all, bot, engine),
functools.partial(message_handlers.delete_all, bot, engine),
commands=["delete_all"],
)
bot.register_message_handler(
functools.partial(handlers.reset_master_pass, bot, engine),
functools.partial(message_handlers.reset_master_pass, bot, engine),
commands=["reset_master_pass"],
)
bot.register_message_handler(
functools.partial(handlers.delete_account, bot, engine),
functools.partial(message_handlers.delete_account, bot, engine),
commands=["delete_account"],
)
bot.register_message_handler(
functools.partial(handlers.help_command, bot),
functools.partial(message_handlers.help_command, bot),
commands=["help", "start"],
)
bot.register_message_handler(
functools.partial(handlers.cancel, bot), commands=["cancel"]
functools.partial(message_handlers.export, bot, engine),
commands=["export"],
)
bot.register_message_handler(
functools.partial(handlers.export, bot, engine), commands=["export"]
)
bot.register_message_handler(
functools.partial(handlers.import_accounts, bot, engine),
functools.partial(message_handlers.import_accounts, bot, engine),
commands=["import"],
)
bot.register_message_handler(
functools.partial(handlers.gen_password, bot),
functools.partial(message_handlers.gen_password, bot),
commands=["gen_password"],
)
bot.register_message_handler(
functools.partial(message_handlers.message_handler, bot),
content_types=[
"text",
"audio",
"document",
"photo",
"sticker",
"video",
"video_note",
"voice",
"location",
"contact",
],
)
bot.register_callback_query_handler(
functools.partial(callback_handlers.delete_message, bot),
lambda call: call.data == "DELETE",
)
return bot

View File

@ -0,0 +1,8 @@
from telebot.async_telebot import AsyncTeleBot
from telebot.types import CallbackQuery as Callback
from . import helper_functions
async def delete_message(bot: AsyncTeleBot, call: Callback) -> None:
await helper_functions.delete_message(bot, call.message)

View File

@ -1,592 +0,0 @@
import functools
import gc
import time
import telebot
from sqlalchemy.future import Engine
from .. import encryption, database, generate_password
from ..account_checks import (
check_account,
check_account_name,
check_login,
check_password,
)
from ..account_parsing import accounts_to_json, json_to_accounts
Message = telebot.types.Message
def _send_tmp_message(
bot: telebot.TeleBot,
chat_id: telebot.types.Message,
text: str,
timeout: int = 5,
) -> None:
bot_mes = bot.send_message(chat_id, text, parse_mode="MarkdownV2")
time.sleep(timeout)
bot.delete_message(chat_id, bot_mes.id)
def _base_handler(
bot: telebot.TeleBot, mes: Message, prev_mes: Message | None = None
) -> None:
bot.delete_message(mes.chat.id, mes.id)
if prev_mes is not None:
bot.delete_message(prev_mes.chat.id, prev_mes.id)
def get_accounts(bot: telebot.TeleBot, engine: Engine, mes: Message) -> None:
_base_handler(bot, mes)
accounts = database.get.get_accounts(engine, mes.from_user.id)
if not accounts:
return _send_tmp_message(bot, mes.chat.id, "У вас нет аккаунтов")
# Make accounts copyable and escape special chars
accounts = [f"`{account}`" for account in accounts]
_send_tmp_message(
bot,
mes.chat.id,
"Ваши аккаунты:\n"
+ "\n".join(accounts)
+ "\nНажмите на название, чтобы скопировать\n"
f"Всего {len(accounts)}",
30,
)
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(
mes, functools.partial(_delete_all, bot, engine, bot_mes)
)
def _delete_all(
bot: telebot.TeleBot, engine: Engine, prev_mes: Message, mes: Message
) -> None:
_base_handler(bot, mes, prev_mes)
text = mes.text.strip()
if text == "YES":
database.delete.purge_accounts(engine, mes.from_user.id)
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, ничего не удалено",
)
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,
"Мастер пароль уже существует",
)
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_handler(bot, mes, prev_mes)
text = mes.text.strip()
if text == "/cancel":
return _send_tmp_message(bot, mes.chat.id, "Успешная отмена")
password_hash, master_salt = encryption.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:
_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, "Мастер пароль не задан")
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_handler(bot, mes, prev_mes)
text = mes.text.strip()
if text == "/cancel":
return _send_tmp_message(bot, mes.chat.id, "Успешная отмена")
hash_, salt = encryption.master_pass.encrypt_master_pass(text)
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, text
gc.collect()
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,
)
if master_password_from_db 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(_add_account2, bot, engine, bot_mes)
)
def _add_account2(
bot: telebot.TeleBot, engine: Engine, prev_mes: Message, mes: Message
) -> None:
_base_handler(bot, mes, prev_mes)
text = mes.text.strip()
if text == "/cancel":
return _send_tmp_message(bot, mes.chat.id, "Успешная отмена")
if not check_account_name(text):
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, "Аккаунт с таким именем уже существует"
)
bot_mes = bot.send_message(mes.chat.id, "Отправьте логин")
data = {"name": text}
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_handler(bot, mes, prev_mes)
text = mes.text.strip()
if text == "/cancel":
return _send_tmp_message(bot, mes.chat.id, "Успешная отмена")
if not check_login(text):
return _send_tmp_message(bot, mes.chat.id, "Не корректный логин")
data["login"] = text
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_handler(bot, mes, prev_mes)
text = mes.text.strip()
if text == "/cancel":
return _send_tmp_message(bot, mes.chat.id, "Успешная отмена")
if not check_password(text):
return _send_tmp_message(bot, mes.chat.id, "Не корректный пароль")
data["passwd"] = text
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_handler(bot, mes, prev_mes)
text = mes.text.strip()
if text == "/cancel":
return _send_tmp_message(bot, mes.chat.id, "Успешная отмена")
salt, hash_ = database.get.get_master_pass(engine, mes.from_user.id)
if not encryption.master_pass.check_master_pass(text, hash_, salt):
return _send_tmp_message(
bot,
mes.chat.id,
"Не подходит главный пароль",
)
name, login, passwd = data["name"], data["login"], data["passwd"]
enc_login, enc_pass, salt = encryption.other_accounts.encrypt(
login,
passwd,
text,
)
result = database.add.add_account(
engine, mes.from_user.id, name, salt, enc_login, enc_pass
)
_send_tmp_message(
bot,
mes.chat.id,
"Успех" if result else "Произошла не предвиденная ошибка",
)
del data, name, login, passwd, enc_login
gc.collect()
def get_account(bot: telebot.TeleBot, engine: Engine, mes: Message) -> None:
_base_handler(bot, mes)
bot_mes = bot.send_message(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, "Нет мастер пароля")
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_handler(bot, mes, prev_mes)
text = mes.text.strip()
if text == "/cancel":
return _send_tmp_message(bot, mes.chat.id, "Успешная отмена")
if text not in database.get.get_accounts(engine, mes.from_user.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(_get_account3, bot, engine, bot_mes, text)
)
def _get_account3(
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,
)
if not encryption.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 = encryption.other_accounts.decrypt(
enc_login,
enc_pass,
text,
salt,
)
_send_tmp_message(
bot,
mes.chat.id,
f"Логин:\n`{login}`\nПароль:\n`{passwd}`\nНажмите на логин "
"или пароль, чтобы скопировать",
30,
)
del text, mes, passwd, login
gc.collect()
def delete_account(bot: telebot.TeleBot, engine: Engine, mes: Message) -> None:
_base_handler(bot, mes)
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, "Нет мастер пароля")
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_handler(bot, mes, prev_mes)
text = mes.text.strip()
if text == "/cancel":
return _send_tmp_message(bot, mes.chat.id, "Успешная отмена")
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, "Аккаунт удалён")
def help_command(bot: telebot.TeleBot, mes: Message) -> None:
message = """Команды:
/set_master_pass - установить мастер пароль
/add_account - создать аккаунт
/get_accounts - получить список аккаунтов
/get_account - получить логин и пароль аккаунта
/delete_account - удалить аккаунт
/delete_all - удалить все аккаунты и мастер пароль
/reset_master_pass - удалить все аккаунты и изменить мастер пароль
/cancel - отмена текущего действия
/help - помощь
/export - получить пароли в json формате
/import - импортировать пароли из json в файле в таком же формате, \
как из /export
/gen_password - создать 10 надёжных паролей"""
bot.send_message(mes.chat.id, message)
def cancel(bot: telebot.TeleBot, mes: Message) -> None:
_base_handler(bot, mes)
_send_tmp_message(bot, mes.chat.id, "Нет активного действия")
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,
)
if master_password_from_db is None:
return _send_tmp_message(bot, mes.chat.id, "Нет мастер пароля")
if not database.get.get_accounts(engine, mes.from_user.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(_export2, bot, engine, bot_mes)
)
def _export2(
bot: telebot.TeleBot, engine: Engine, prev_mes: Message, 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 not encryption.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)
accounts = encryption.other_accounts.decrypt_multiple(accounts, text)
json_io = accounts_to_json(accounts)
bot_mes = bot.send_document(mes.chat.id, json_io)
del text, accounts, json_io
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:
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, "Файл слишком большой")
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:
return _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 not encryption.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 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 = encryption.other_accounts.encrypt(
login,
passwd,
text,
)
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()
def gen_password(bot: telebot.TeleBot, mes: Message) -> None:
_base_handler(bot, mes)
# Generate 10 passwords and put 'em in the backticks
passwords = (f"`{generate_password.gen_password()}`" for _ in range(10))
text = (
"Пароли:\n"
+ "\n".join(passwords)
+ "\nНажмите на пароль, чтобы его скопировать"
)
_send_tmp_message(bot, mes.chat.id, text, 15)

View File

@ -0,0 +1,95 @@
import asyncio
from typing import Any, Awaitable, Callable
import telebot
from telebot.async_telebot import AsyncTeleBot
from . import markups
Message = telebot.types.Message
Handler = Callable[[Message], Awaitable[Any]]
Markups = (
telebot.types.ReplyKeyboardMarkup
| telebot.types.InlineKeyboardMarkup
| telebot.types.ReplyKeyboardRemove
)
states: dict[tuple[int, int], Handler] = {}
def register_state(
message: Message,
handler: Callable[[Message], Any],
) -> None:
states[(message.chat.id, message.from_user.id)] = handler
def reset_state(message: Message) -> None:
try:
del states[(message.chat.id, message.from_user.id)]
except KeyError:
pass
def get_state(message: Message) -> Handler | None:
return states.get((message.chat.id, message.from_user.id))
async def delete_message(
bot: AsyncTeleBot,
mes: Message,
*,
sleep_time: int = 0,
) -> bool:
try:
if sleep_time != 0:
await asyncio.sleep(sleep_time)
await bot.delete_message(mes.chat.id, mes.id)
except telebot.apihelper.ApiException:
return False
else:
return True
async def send_tmp_message(
bot: AsyncTeleBot,
chat_id: telebot.types.Message,
text: str,
*,
sleep_time: int = 5,
markup: Markups | None = None,
) -> None:
bot_mes = await bot.send_message(
chat_id, text, parse_mode="MarkdownV2", reply_markup=markup
)
await delete_message(bot, bot_mes, sleep_time=sleep_time)
async def base_handler(
bot: AsyncTeleBot, mes: Message, prev_mes: Message | None = None
) -> None:
reset_state(mes)
await delete_message(bot, mes)
if prev_mes is not None:
await delete_message(bot, prev_mes)
async def send_deleteable_message(
bot: AsyncTeleBot,
chat_id: int,
text: str,
) -> None:
"""Sends a message with a delete button"""
markup = markups.deletion_markup()
await bot.send_message(
chat_id,
text,
parse_mode="MarkdownV2",
reply_markup=markup,
)
def escape(text: str) -> str:
escaped_chars = "*_~|`[("
for char in escaped_chars:
text = text.replace(char, rf"\\{char}")
return text

18
src/bot/markups.py Normal file
View File

@ -0,0 +1,18 @@
from telebot.types import (
InlineKeyboardButton,
InlineKeyboardMarkup,
ReplyKeyboardMarkup,
)
def deletion_markup() -> InlineKeyboardMarkup:
markup = InlineKeyboardMarkup(row_width=5)
button = InlineKeyboardButton("Удалить сообщение", callback_data="DELETE")
markup.add(button)
return markup
def account_markup(account_names: list[str]) -> ReplyKeyboardMarkup:
markup = ReplyKeyboardMarkup(resize_keyboard=True, row_width=2)
markup.add(*account_names)
return markup

683
src/bot/message_handlers.py Normal file
View File

@ -0,0 +1,683 @@
import functools
import gc
import telebot
from sqlalchemy.future import Engine
from telebot.async_telebot import AsyncTeleBot
from .. import database, encryption, generate_password
from ..account_checks import (
check_account,
check_account_name,
check_login,
check_password,
)
from ..account_parsing import accounts_to_json, json_to_accounts
from ..decrypted_account import DecryptedAccount
from . import markups
from .helper_functions import (
base_handler,
delete_message,
escape,
get_state,
register_state,
send_deleteable_message,
send_tmp_message,
)
Message = telebot.types.Message
async def get_accounts(
bot: AsyncTeleBot,
engine: Engine,
mes: Message,
) -> None:
await base_handler(bot, mes)
accounts = database.get.get_accounts(
engine,
mes.from_user.id,
to_sort=True,
)
if not accounts:
return await send_tmp_message(bot, mes.chat.id, "У вас нет аккаунтов")
# Make accounts copyable and escape special chars
accounts = [f"`{account}`" for account in accounts]
await send_deleteable_message(
bot,
mes.chat.id,
"Ваши аккаунты:\n"
+ "\n".join(accounts)
+ "\nНажмите на название, чтобы скопировать\n"
f"Всего {len(accounts)}",
)
async def delete_all(bot: AsyncTeleBot, engine: Engine, mes: Message) -> None:
await base_handler(bot, mes)
bot_mes = await bot.send_message(
mes.chat.id,
"Вы действительно хотите удалить все ваши аккаунты? Это действие "
"нельзя отменить. "
"Отправьте YES для подтверждения",
)
register_state(mes, functools.partial(_delete_all2, bot, engine, bot_mes))
async def _delete_all2(
bot: AsyncTeleBot, engine: Engine, prev_mes: Message, mes: Message
) -> None:
await base_handler(bot, mes, prev_mes)
text = mes.text.strip()
if text == "YES":
database.delete.purge_accounts(engine, mes.from_user.id)
database.delete.delete_master_pass(engine, mes.from_user.id)
await send_tmp_message(
bot,
mes.chat.id,
"Всё успешно удалено",
sleep_time=10,
)
else:
await send_tmp_message(
bot,
mes.chat.id,
"Вы отправили не YES, ничего не удалено",
)
async def set_master_password(
bot: AsyncTeleBot,
engine: Engine,
mes: Message,
) -> None:
await base_handler(bot, mes, None)
if database.get.get_master_pass(engine, mes.from_user.id) is not None:
return await send_tmp_message(
bot,
mes.chat.id,
"Мастер пароль уже существует",
)
bot_mes = await bot.send_message(mes.chat.id, "Отправьте мастер пароль")
register_state(
mes,
functools.partial(_set_master_pass2, bot, engine, bot_mes),
)
async def _set_master_pass2(
bot: AsyncTeleBot, engine: Engine, prev_mes: Message, mes: Message
) -> None:
await base_handler(bot, mes, prev_mes)
text = mes.text.strip()
if text == "/cancel":
return await send_tmp_message(bot, mes.chat.id, "Успешная отмена")
master_password = encryption.master_pass.encrypt_master_pass(
mes.from_user.id,
text,
)
database.add.add_master_pass(engine, master_password)
await send_tmp_message(bot, mes.chat.id, "Успех")
del mes, text
gc.collect()
async def reset_master_pass(
bot: AsyncTeleBot,
engine: Engine,
mes: Message,
) -> None:
await base_handler(bot, mes)
if database.get.get_master_pass(engine, mes.from_user.id) is None:
return await send_tmp_message(
bot,
mes.chat.id,
"Мастер пароль не задан",
)
bot_mes = await bot.send_message(
mes.chat.id,
"Отправьте новый мастер пароль, осторожно, все текущие аккаунты "
"будут удалены навсегда",
)
register_state(
mes,
functools.partial(_reset_master_pass2, bot, engine, bot_mes),
)
async def _reset_master_pass2(
bot: AsyncTeleBot, engine: Engine, prev_mes: Message, mes: Message
) -> None:
await base_handler(bot, mes, prev_mes)
text = mes.text.strip()
if text == "/cancel":
return await send_tmp_message(bot, mes.chat.id, "Успешная отмена")
master_password = encryption.master_pass.encrypt_master_pass(
mes.from_user.id,
text,
)
database.delete.purge_accounts(engine, mes.from_user.id)
database.change.change_master_pass(engine, master_password)
await send_tmp_message(
bot, mes.chat.id, "Все ваши аккаунты удалены, а мастер пароль изменён"
)
del mes, text
gc.collect()
async def add_account(bot: AsyncTeleBot, engine: Engine, mes: Message) -> None:
await 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 await send_tmp_message(bot, mes.chat.id, "Нет мастер пароля")
bot_mes = await bot.send_message(
mes.chat.id,
"Отправьте название аккаунта",
)
register_state(
mes,
functools.partial(_add_account2, bot, engine, bot_mes),
)
async def _add_account2(
bot: AsyncTeleBot, engine: Engine, prev_mes: Message, mes: Message
) -> None:
await base_handler(bot, mes, prev_mes)
text = mes.text.strip()
if text == "/cancel":
return await send_tmp_message(bot, mes.chat.id, "Успешная отмена")
if not check_account_name(text):
return await send_tmp_message(
bot,
mes.chat.id,
"Не корректное название аккаунта",
)
if text in database.get.get_accounts(engine, mes.from_user.id):
return await send_tmp_message(
bot, mes.chat.id, "Аккаунт с таким именем уже существует"
)
bot_mes = await bot.send_message(mes.chat.id, "Отправьте логин")
data = {"name": text}
register_state(
mes,
functools.partial(_add_account3, bot, engine, bot_mes, data),
)
async def _add_account3(
bot: AsyncTeleBot,
engine: Engine,
prev_mes: Message,
data: dict[str, str],
mes: Message,
) -> None:
await base_handler(bot, mes, prev_mes)
text = mes.text.strip()
if text == "/cancel":
return await send_tmp_message(bot, mes.chat.id, "Успешная отмена")
if not check_login(text):
return await send_tmp_message(bot, mes.chat.id, "Не корректный логин")
data["login"] = text
bot_mes = await bot.send_message(
mes.chat.id,
"Отправьте пароль от аккаунта",
)
register_state(
mes,
functools.partial(_add_account4, bot, engine, bot_mes, data),
)
async def _add_account4(
bot: AsyncTeleBot,
engine: Engine,
prev_mes: Message,
data: dict[str, str],
mes: Message,
) -> None:
await base_handler(bot, mes, prev_mes)
text = mes.text.strip()
if text == "/cancel":
return await send_tmp_message(bot, mes.chat.id, "Успешная отмена")
if not check_password(text):
return await send_tmp_message(
bot,
mes.chat.id,
"Не корректный пароль",
)
data["passwd"] = text
bot_mes = await bot.send_message(mes.chat.id, "Отправьте мастер пароль")
register_state(
mes,
functools.partial(_add_account5, bot, engine, bot_mes, data),
)
async def _add_account5(
bot: AsyncTeleBot,
engine: Engine,
prev_mes: Message,
data: dict[str, str],
mes: Message,
) -> None:
await base_handler(bot, mes, prev_mes)
text = mes.text.strip()
if text == "/cancel":
return await send_tmp_message(bot, mes.chat.id, "Успешная отмена")
master_password = database.get.get_master_pass(engine, mes.from_user.id)
if not encryption.master_pass.check_master_pass(text, master_password):
return await send_tmp_message(
bot,
mes.chat.id,
"Не подходит главный пароль",
)
# 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"],
)
encrypted_account = encryption.other_accounts.encrypt(
account,
text,
)
result = database.add.add_account(
engine,
encrypted_account,
)
await send_tmp_message(
bot,
mes.chat.id,
"Успех" if result else "Произошла не предвиденная ошибка",
)
del data, account
gc.collect()
async def get_account(bot: AsyncTeleBot, engine: Engine, mes: Message) -> None:
await base_handler(bot, mes)
master_pass = database.get.get_master_pass(engine, mes.from_user.id)
if master_pass is None:
return await send_tmp_message(bot, mes.chat.id, "Нет мастер пароля")
accounts = database.get.get_accounts(
engine,
mes.from_user.id,
to_sort=True,
)
markup = markups.account_markup(accounts)
bot_mes = await bot.send_message(
mes.chat.id, "Отправьте название аккаунта", reply_markup=markup
)
register_state(
mes,
functools.partial(_get_account2, bot, engine, bot_mes),
)
async def _get_account2(
bot: AsyncTeleBot, engine: Engine, prev_mes: Message, mes: Message
) -> None:
await base_handler(bot, mes, prev_mes)
text = mes.text.strip()
if text == "/cancel":
return await send_tmp_message(bot, mes.chat.id, "Успешная отмена")
if text not in database.get.get_accounts(engine, mes.from_user.id):
return await send_tmp_message(bot, mes.chat.id, "Нет такого аккаунта")
bot_mes = await bot.send_message(mes.chat.id, "Отправьте мастер пароль")
register_state(
mes,
functools.partial(_get_account3, bot, engine, bot_mes, text),
)
async def _get_account3(
bot: AsyncTeleBot,
engine: Engine,
prev_mes: Message,
name: str,
mes: Message,
) -> None:
await base_handler(bot, mes, prev_mes)
text = mes.text.strip()
if text == "/cancel":
return await send_tmp_message(bot, mes.chat.id, "Успешная отмена")
master_password = database.get.get_master_pass(
engine,
mes.from_user.id,
)
if not encryption.master_pass.check_master_pass(text, master_password):
return await send_tmp_message(
bot,
mes.chat.id,
"Не подходит мастер пароль",
)
account = database.get.get_account_info(engine, mes.from_user.id, name)
account = encryption.other_accounts.decrypt(
account,
text,
)
await send_deleteable_message(
bot,
mes.chat.id,
f"Название:\n{escape(account.name)}\n"
f"Логин:\n`{account.login}`\nПароль:\n`{account.password}`\nНажмите "
"на логин или пароль, чтобы скопировать",
)
del text, mes
gc.collect()
async def delete_account(
bot: AsyncTeleBot,
engine: Engine,
mes: Message,
) -> None:
await base_handler(bot, mes)
master_pass = database.get.get_master_pass(engine, mes.from_user.id)
if master_pass is None:
return await send_tmp_message(bot, mes.chat.id, "Нет мастер пароля")
accounts = database.get.get_accounts(
engine,
mes.from_user.id,
to_sort=True,
)
markup = markups.account_markup(accounts)
bot_mes = await bot.send_message(
mes.chat.id,
"Отправьте название аккаунта, который вы хотите удалить",
reply_markup=markup,
)
register_state(
mes,
functools.partial(_delete_account2, bot, engine, bot_mes),
)
async def _delete_account2(
bot: AsyncTeleBot,
engine: Engine,
prev_mes: Message,
mes: Message,
):
await base_handler(bot, mes, prev_mes)
text = mes.text.strip()
if text == "/cancel":
return await send_tmp_message(bot, mes.chat.id, "Успешная отмена")
if text not in database.get.get_accounts(engine, mes.from_user.id):
return await send_tmp_message(bot, mes.chat.id, "Нет такого аккаунта")
bot_mes = await bot.send_message(
mes.from_user.id,
f'Вы уверены, что хотите удалить аккаунт "{text}"?\nОтправьте YES для '
"подтверждения",
)
register_state(
mes,
functools.partial(_delete_account3, bot, engine, bot_mes, text),
)
async def _delete_account3(
bot: AsyncTeleBot,
engine: Engine,
prev_mes: Message,
account_name: str,
mes: Message,
) -> None:
await base_handler(bot, mes, prev_mes)
text = mes.text.strip()
if text != "YES":
return await send_tmp_message(bot, mes.chat.id, "Успешная отмена")
database.delete.delete_account(engine, mes.from_user.id, account_name)
await send_tmp_message(bot, mes.chat.id, "Аккаунт удалён")
async def help_command(bot: AsyncTeleBot, mes: Message) -> None:
message = """Команды:
/set_master_pass - установить мастер пароль
/add_account - создать аккаунт
/get_accounts - получить список аккаунтов
/get_account - получить логин и пароль аккаунта
/delete_account - удалить аккаунт
/delete_all - удалить все аккаунты и мастер пароль
/reset_master_pass - удалить все аккаунты и изменить мастер пароль
/cancel - отмена текущего действия
/help - помощь
/export - получить пароли в json формате
/import - импортировать пароли из json в файле в таком же формате, \
как из /export
/gen_password - создать 10 надёжных паролей"""
await bot.send_message(mes.chat.id, message)
async def export(bot: AsyncTeleBot, engine: Engine, mes: Message) -> None:
await 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 await send_tmp_message(bot, mes.chat.id, "Нет мастер пароля")
if not database.get.get_accounts(engine, mes.from_user.id):
return await send_tmp_message(bot, mes.chat.id, "Нет аккаунтов")
bot_mes = await bot.send_message(mes.chat.id, "Отправьте мастер пароль")
register_state(mes, functools.partial(_export2, bot, engine, bot_mes))
async def _export2(
bot: AsyncTeleBot, engine: Engine, prev_mes: Message, mes: Message
) -> None:
await base_handler(bot, mes, prev_mes)
text = mes.text.strip()
if text == "/cancel":
return await send_tmp_message(bot, mes.chat.id, "Успешная отмена")
master_password = database.get.get_master_pass(
engine,
mes.from_user.id,
)
if not encryption.master_pass.check_master_pass(text, master_password):
return await send_tmp_message(
bot,
mes.chat.id,
"Не подходит мастер пароль",
)
accounts = database.get.get_all_accounts(engine, mes.from_user.id)
accounts = encryption.other_accounts.decrypt_multiple(accounts, text)
json_io = accounts_to_json(accounts)
await bot.send_document(
mes.chat.id,
json_io,
reply_markup=markups.deletion_markup(),
)
del text, accounts, json_io
gc.collect()
async def import_accounts(
bot: AsyncTeleBot,
engine: Engine,
mes: Message,
) -> None:
await 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 await send_tmp_message(bot, mes.chat.id, "Нет мастер пароля")
bot_mes = await bot.send_message(mes.chat.id, "Отправьте json файл")
register_state(mes, functools.partial(_import2, bot, engine, bot_mes))
async def _import2(
bot: AsyncTeleBot, engine: Engine, prev_mes: Message, mes: Message
) -> None:
await base_handler(bot, mes, prev_mes)
if mes.text is not None:
text = mes.text.strip()
if text == "/cancel":
return await send_tmp_message(bot, mes.chat.id, "Успешная отмена")
if mes.document is None:
return await send_tmp_message(
bot,
mes.chat.id,
"Вы должны отправить документ",
)
if mes.document.file_size > 102_400: # If file size is bigger than 100 MB
return await send_tmp_message(
bot,
mes.chat.id,
"Файл слишком большой",
)
file_info = await bot.get_file(mes.document.file_id)
downloaded_file = await bot.download_file(file_info.file_path)
try:
accounts = json_to_accounts(
downloaded_file.decode("utf-8"),
mes.from_user.id,
)
except Exception:
return await send_tmp_message(
bot,
mes.chat.id,
"Ошибка во время работы с файлом",
)
bot_mes = await bot.send_message(mes.chat.id, "Отправьте мастер пароль")
register_state(
mes,
functools.partial(_import3, bot, engine, bot_mes, accounts),
)
async def _import3(
bot: AsyncTeleBot,
engine: Engine,
prev_mes: Message,
accounts: list[DecryptedAccount],
mes: Message,
) -> None:
await base_handler(bot, mes, prev_mes)
text = mes.text.strip()
if text == "/cancel":
return await send_tmp_message(bot, mes.chat.id, "Успешная отмена")
master_password = database.get.get_master_pass(
engine,
mes.from_user.id,
)
if not encryption.master_pass.check_master_pass(text, master_password):
return await send_tmp_message(
bot,
mes.chat.id,
"Не подходит мастер пароль",
)
# List of names of accounts, which failed to be added to the database
# or failed the tests
failed: list[str] = []
for account in accounts:
if not check_account(account):
failed.append(account.name)
continue
account = encryption.other_accounts.encrypt(
account,
text,
)
result = database.add.add_account(engine, account)
if not result:
failed.append(account.name)
if failed:
mes_text = "Не удалось добавить:\n" + "\n".join(failed)
else:
mes_text = "Успех"
await send_tmp_message(bot, mes.chat.id, mes_text, 10)
del text, mes, accounts
gc.collect()
async def gen_password(bot: AsyncTeleBot, mes: Message) -> None:
await base_handler(bot, mes)
# Generate 10 passwords and put them in the backticks
passwords = (f"`{generate_password.gen_password()}`" for _ in range(10))
text = (
"Пароли:\n"
+ "\n".join(passwords)
+ "\nНажмите на пароль, чтобы его скопировать"
)
await send_deleteable_message(bot, mes.chat.id, text)
async def message_handler(bot: AsyncTeleBot, mes: Message) -> None:
handler = get_state(mes)
if handler is None:
await delete_message(bot, mes)
if mes.text.strip() == "/cancel":
await send_tmp_message(bot, mes.chat.id, "Нет активного действия")
await send_tmp_message(
bot,
mes.chat.id,
"Вы отправили не корректное сообщение",
)
return
await handler(mes)

View File

@ -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"]

View File

@ -5,23 +5,9 @@ from sqlalchemy.future import Engine
from . import models
def add_account(
engine: Engine,
user_id: int,
name: str,
salt: bytes,
enc_login: bytes,
enc_password: bytes,
) -> bool:
def add_account(engine: Engine, account: models.Account) -> bool:
"""Adds account to the database. Returns true on success,
false otherwise"""
account = models.Account(
user_id=user_id,
name=name,
salt=salt,
enc_login=enc_login,
enc_password=enc_password,
)
try:
with sqlmodel.Session(engine) as session:
session.add(account)
@ -32,19 +18,9 @@ def add_account(
return True
def add_master_pass(
engine: Engine,
user_id: int,
salt: bytes,
password_hash: bytes,
) -> bool:
def add_master_pass(engine: Engine, master_pass: models.MasterPass) -> bool:
"""Adds master password the database. Returns true on success,
false otherwise"""
master_pass = models.MasterPass(
user_id=user_id,
salt=salt,
password_hash=password_hash,
)
try:
with sqlmodel.Session(engine) as session:
session.add(master_pass)

View File

@ -5,13 +5,17 @@ from . import models
def change_master_pass(
engine: Engine, user_id: int, salt: bytes, password: bytes
engine: Engine,
master_password: models.MasterPass,
) -> None:
"""Changes master password and salt in the database"""
statement = (
sqlmodel.update(models.MasterPass)
.where(models.MasterPass.user_id == user_id)
.values(salt=salt, password_hash=password)
.where(models.MasterPass.user_id == master_password.user_id)
.values(
salt=master_password.salt,
password_hash=master_password.password_hash,
)
)
with sqlmodel.Session(engine) as session:
session.exec(statement)

View File

@ -1,5 +1,3 @@
from typing import Iterator
import sqlmodel
from sqlalchemy.future import Engine
@ -9,21 +7,35 @@ from . import models
def get_master_pass(
engine: Engine,
user_id: int,
) -> tuple[bytes, bytes] | None:
"""Gets master pass. Returns tuple of salt and password
or None if it wasn't found"""
) -> models.MasterPass | None:
"""Gets master password of a user"""
statement = sqlmodel.select(models.MasterPass).where(
models.MasterPass.user_id == user_id,
)
with sqlmodel.Session(engine) as session:
result = session.exec(statement).first()
if result is None:
return
return (result.salt, result.password_hash)
return result
def get_accounts(engine: Engine, user_id: int) -> list[str]:
"""Gets list of account names"""
def get_accounts(
engine: Engine,
user_id: int,
*,
to_sort: bool = False,
) -> list[str]:
"""Gets a list of account names of a user"""
statement = sqlmodel.select(models.Account.name).where(
models.Account.user_id == user_id,
)
if to_sort:
statement = statement.order_by(models.Account.name)
with sqlmodel.Session(engine) as session:
result = session.exec(statement).fetchall()
return result
def get_all_accounts(engine: Engine, user_id: int) -> list[models.Account]:
"""Returns a list of accounts of a user"""
statement = (
sqlmodel.select(models.Account)
.where(
@ -32,46 +44,20 @@ def get_accounts(engine: Engine, user_id: int) -> list[str]:
.order_by(models.Account.name)
)
with sqlmodel.Session(engine) as session:
result = session.exec(statement)
return [account.name for account in result]
def get_all_accounts(
engine: Engine, user_id: int
) -> Iterator[tuple[str, bytes, bytes, bytes]]:
"""Returns an iterator of tuples, where values represent account's
name, salt, encrypted login and encrypted password"""
statement = (
sqlmodel.select(models.Account)
.where(
models.Account.user_id == user_id,
)
.order_by(models.Account.name)
)
with sqlmodel.Session(engine) as session:
result = session.exec(statement)
yield from (
(
account.name,
account.salt,
account.enc_login,
account.enc_password,
)
for account in result
)
result = session.exec(statement).fetchall()
return result
def get_account_info(
engine: Engine, user_id: int, name: str
) -> tuple[bytes, bytes, bytes]:
"""Gets account info. Returns tuple of salt, login and password
or None if it wasn't found"""
engine: Engine,
user_id: int,
name: str,
) -> models.Account:
"""Gets account info"""
statement = sqlmodel.select(models.Account).where(
models.Account.user_id == user_id,
models.Account.name == name,
)
with sqlmodel.Session(engine) as session:
result = session.exec(statement).first()
if result is None:
return
return (result.salt, result.enc_login, result.enc_password)
return result

8
src/decrypted_account.py Normal file
View File

@ -0,0 +1,8 @@
import pydantic
class DecryptedAccount(pydantic.BaseModel):
user_id: int
name: str
login: str
password: str

View File

@ -3,6 +3,8 @@ import os
from cryptography.exceptions import InvalidKey
from cryptography.hazmat.primitives.kdf.scrypt import Scrypt
from ..database.models import MasterPass
MEMORY_USAGE = 2**14
@ -17,22 +19,23 @@ def _get_kdf(salt: bytes) -> Scrypt:
return kdf
def encrypt_master_pass(password: str) -> tuple[bytes, bytes]:
"""Hashes master password and return tuple of hashed password and salt"""
def encrypt_master_pass(user_id: int, password: str) -> MasterPass:
"""Hashes master password and returns MasterPass object"""
salt = os.urandom(64)
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(
password: str,
enc_password: bytes,
salt: bytes,
) -> bool:
def check_master_pass(password: str, master_password: MasterPass) -> bool:
"""Checks if the master password is correct"""
kdf = _get_kdf(salt)
kdf = _get_kdf(master_password.salt)
try:
kdf.verify(password.encode("utf-8"), enc_password)
kdf.verify(password.encode("utf-8"), master_password.password_hash)
except InvalidKey:
return False
else:

View File

@ -1,14 +1,18 @@
import base64
import os
from typing import Iterator
from typing import Iterable, Iterator
from cryptography.fernet import Fernet
from cryptography.hazmat.backends import default_backend
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
from ..database.models import Account
from ..decrypted_account import DecryptedAccount
def _generate_key(salt: bytes, master_pass: bytes) -> bytes:
"""Generates key for fernet encryption"""
kdf = PBKDF2HMAC(
algorithm=hashes.SHA256(),
length=32,
@ -21,49 +25,58 @@ def _generate_key(salt: bytes, master_pass: bytes) -> bytes:
def encrypt(
login: str,
passwd: str,
account: DecryptedAccount,
master_pass: str,
) -> tuple[bytes, bytes, bytes]:
"""Encrypts login and password of a user using their master
password as a key.
Returns a tuple of encrypted login, password and salt"""
) -> Account:
"""Encrypts account using master password and returns Account object"""
salt = os.urandom(64)
key = _generate_key(salt, master_pass.encode("utf-8"))
f = Fernet(key)
enc_login = base64.urlsafe_b64decode(f.encrypt(login.encode("utf-8")))
enc_password = base64.urlsafe_b64decode(f.encrypt(passwd.encode("utf-8")))
return (enc_login, enc_password, salt)
enc_login = base64.urlsafe_b64decode(
f.encrypt(account.login.encode("utf-8")),
)
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(
enc_login: bytes,
enc_pass: bytes,
account: Account,
master_pass: str,
salt: bytes,
) -> tuple[str, str]:
"""Decrypts login and password using their
master password as a key.
Returns a tuple of decrypted login and password"""
key = _generate_key(salt, master_pass.encode("utf-8"))
) -> DecryptedAccount:
"""Decrypts account using master password and returns
DecryptedAccount object"""
key = _generate_key(account.salt, master_pass.encode("utf-8"))
f = Fernet(key)
login = f.decrypt(base64.urlsafe_b64encode(enc_login)).decode("utf-8")
password = f.decrypt(base64.urlsafe_b64encode(enc_pass)).decode("utf-8")
return (login, password)
login = f.decrypt(
base64.urlsafe_b64encode(account.enc_login),
).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(
accounts: Iterator[tuple[str, bytes, bytes, bytes]], master_pass: str
) -> Iterator[tuple[str, str, str]]:
"""Gets an iterator of tuples, where values represent account's name, salt,
encrypted login and encrypted password.
Return an iterator of names, logins and passwords as a tuple"""
accounts: Iterable[Account], master_pass: str
) -> Iterator[DecryptedAccount]:
"""Decrypts an iterable of accounts using master_pass and
returns an Iterator of decrypted accounts"""
for account in accounts:
name, salt, enc_login, enc_passwd = account
login, passwd = decrypt(
enc_login,
enc_passwd,
master_pass,
salt,
)
yield (name, login, passwd)
yield decrypt(account, master_pass)