Compare commits

..

6 Commits

Author SHA1 Message Date
fdbed91512 Now account name is copyable 2023-01-03 12:12:40 +03:00
f9d163361a Renamed functions in db.get
get_accounts -> get_account_names
get_all_accounts -> get_accounts
2023-01-03 12:08:06 +03:00
9ec66a3521 Renamed src/database into src/db 2023-01-03 11:59:06 +03:00
70e9afe21d Small reformating in bot.message_handler 2023-01-03 11:48:25 +03:00
5d59adb7d2 Renamed encryption/other_accounts into encryption/accounts 2023-01-01 00:21:53 +03:00
281c4a262b _export2 no longer blocks event loop during decryption
Removed sorting in get_all_accounts
Removed decrypt_multiple function because it is no longer used
Now running decrytion of accounts in ProcessPoolExecutor
2023-01-01 00:18:57 +03:00
13 changed files with 60 additions and 78 deletions

View File

@ -7,7 +7,7 @@ from . import (
account_checks, account_checks,
account_parsing, account_parsing,
bot, bot,
database, db,
decrypted_account, decrypted_account,
encryption, encryption,
generate_password, generate_password,
@ -19,19 +19,19 @@ __all__ = [
"bot", "bot",
"decrypted_account", "decrypted_account",
"encryption", "encryption",
"database", "db",
"generate_password", "generate_password",
] ]
def main() -> None: def main() -> None:
load_dotenv("./.env") load_dotenv("./.env")
engine = database.prepare.get_engine( engine = db.prepare.get_engine(
host=os.getenv("DB_HOST"), host=os.getenv("DB_HOST"),
user=os.getenv("DB_USER"), user=os.getenv("DB_USER"),
passwd=os.getenv("DB_PASS"), passwd=os.getenv("DB_PASS"),
db=os.getenv("DB_NAME"), db=os.getenv("DB_NAME"),
) )
database.prepare.prepare(engine) db.prepare.prepare(engine)
bot_ = bot.create_bot(os.getenv("TG_TOKEN"), engine) bot_ = bot.create_bot(os.getenv("TG_TOKEN"), engine)
asyncio.run(bot_.infinity_polling()) asyncio.run(bot_.infinity_polling())

View File

@ -86,10 +86,3 @@ async def send_deleteable_message(
parse_mode="MarkdownV2", parse_mode="MarkdownV2",
reply_markup=markup, reply_markup=markup,
) )
def escape(text: str) -> str:
escaped_chars = "*_~|`[("
for char in escaped_chars:
text = text.replace(char, rf"\\{char}")
return text

View File

@ -1,11 +1,13 @@
import asyncio
import functools import functools
import gc import gc
from concurrent.futures import ProcessPoolExecutor
import telebot import telebot
from sqlalchemy.future import Engine from sqlalchemy.future import Engine
from telebot.async_telebot import AsyncTeleBot from telebot.async_telebot import AsyncTeleBot
from .. import database, encryption, generate_password from .. import db, encryption, generate_password
from ..account_checks import ( from ..account_checks import (
check_account, check_account,
check_account_name, check_account_name,
@ -18,7 +20,6 @@ from . import markups
from .helper_functions import ( from .helper_functions import (
base_handler, base_handler,
delete_message, delete_message,
escape,
get_state, get_state,
register_state, register_state,
send_deleteable_message, send_deleteable_message,
@ -34,7 +35,7 @@ async def get_accounts(
mes: Message, mes: Message,
) -> None: ) -> None:
await base_handler(bot, mes) await base_handler(bot, mes)
accounts = database.get.get_accounts( accounts = db.get.get_account_names(
engine, engine,
mes.from_user.id, mes.from_user.id,
to_sort=True, to_sort=True,
@ -71,8 +72,8 @@ async def _delete_all2(
await base_handler(bot, mes, prev_mes) await base_handler(bot, mes, prev_mes)
text = mes.text.strip() text = mes.text.strip()
if text == "YES": if text == "YES":
database.delete.purge_accounts(engine, mes.from_user.id) db.delete.purge_accounts(engine, mes.from_user.id)
database.delete.delete_master_pass(engine, mes.from_user.id) db.delete.delete_master_pass(engine, mes.from_user.id)
await send_tmp_message( await send_tmp_message(
bot, bot,
mes.chat.id, mes.chat.id,
@ -93,7 +94,7 @@ async def set_master_password(
mes: Message, mes: Message,
) -> None: ) -> None:
await base_handler(bot, mes, None) await base_handler(bot, mes, None)
if database.get.get_master_pass(engine, mes.from_user.id) is not None: if db.get.get_master_pass(engine, mes.from_user.id) is not None:
return await send_tmp_message( return await send_tmp_message(
bot, bot,
mes.chat.id, mes.chat.id,
@ -118,7 +119,7 @@ async def _set_master_pass2(
mes.from_user.id, mes.from_user.id,
text, text,
) )
database.add.add_master_pass(engine, master_password) db.add.add_master_pass(engine, master_password)
await send_tmp_message(bot, mes.chat.id, "Успех") await send_tmp_message(bot, mes.chat.id, "Успех")
del mes, text del mes, text
@ -132,7 +133,7 @@ async def reset_master_pass(
) -> None: ) -> None:
await base_handler(bot, mes) await base_handler(bot, mes)
if database.get.get_master_pass(engine, mes.from_user.id) is None: if db.get.get_master_pass(engine, mes.from_user.id) is None:
return await send_tmp_message( return await send_tmp_message(
bot, bot,
mes.chat.id, mes.chat.id,
@ -163,8 +164,8 @@ async def _reset_master_pass2(
mes.from_user.id, mes.from_user.id,
text, text,
) )
database.delete.purge_accounts(engine, mes.from_user.id) db.delete.purge_accounts(engine, mes.from_user.id)
database.change.change_master_pass(engine, master_password) db.change.change_master_pass(engine, master_password)
await send_tmp_message( await send_tmp_message(
bot, mes.chat.id, "Все ваши аккаунты удалены, а мастер пароль изменён" bot, mes.chat.id, "Все ваши аккаунты удалены, а мастер пароль изменён"
@ -176,7 +177,7 @@ async def _reset_master_pass2(
async def add_account(bot: AsyncTeleBot, engine: Engine, mes: Message) -> None: async def add_account(bot: AsyncTeleBot, engine: Engine, mes: Message) -> None:
await base_handler(bot, mes) await base_handler(bot, mes)
master_password_from_db = database.get.get_master_pass( master_password_from_db = db.get.get_master_pass(
engine, engine,
mes.from_user.id, mes.from_user.id,
) )
@ -208,7 +209,7 @@ async def _add_account2(
mes.chat.id, mes.chat.id,
"Не корректное название аккаунта", "Не корректное название аккаунта",
) )
if text in database.get.get_accounts(engine, mes.from_user.id): if text in db.get.get_account_names(engine, mes.from_user.id):
return await send_tmp_message( return await send_tmp_message(
bot, mes.chat.id, "Аккаунт с таким именем уже существует" bot, mes.chat.id, "Аккаунт с таким именем уже существует"
) )
@ -289,7 +290,7 @@ async def _add_account5(
if text == "/cancel": if text == "/cancel":
return await send_tmp_message(bot, mes.chat.id, "Успешная отмена") return await send_tmp_message(bot, mes.chat.id, "Успешная отмена")
master_password = database.get.get_master_pass(engine, mes.from_user.id) master_password = db.get.get_master_pass(engine, mes.from_user.id)
if not encryption.master_pass.check_master_pass(text, master_password): if not encryption.master_pass.check_master_pass(text, master_password):
return await send_tmp_message( return await send_tmp_message(
bot, bot,
@ -305,12 +306,9 @@ async def _add_account5(
password=data["passwd"], password=data["passwd"],
) )
encrypted_account = encryption.other_accounts.encrypt( encrypted_account = encryption.accounts.encrypt(account, text)
account,
text,
)
result = database.add.add_account( result = db.add.add_account(
engine, engine,
encrypted_account, encrypted_account,
) )
@ -329,11 +327,11 @@ async def _add_account5(
async def get_account(bot: AsyncTeleBot, engine: Engine, mes: Message) -> None: async def get_account(bot: AsyncTeleBot, engine: Engine, mes: Message) -> None:
await base_handler(bot, mes) await base_handler(bot, mes)
master_pass = database.get.get_master_pass(engine, mes.from_user.id) master_pass = db.get.get_master_pass(engine, mes.from_user.id)
if master_pass is None: if master_pass is None:
return await send_tmp_message(bot, mes.chat.id, "Нет мастер пароля") return await send_tmp_message(bot, mes.chat.id, "Нет мастер пароля")
accounts = database.get.get_accounts( accounts = db.get.get_account_names(
engine, engine,
mes.from_user.id, mes.from_user.id,
to_sort=True, to_sort=True,
@ -356,7 +354,7 @@ async def _get_account2(
if text == "/cancel": if text == "/cancel":
return await send_tmp_message(bot, mes.chat.id, "Успешная отмена") return await send_tmp_message(bot, mes.chat.id, "Успешная отмена")
if text not in database.get.get_accounts(engine, mes.from_user.id): if text not in db.get.get_account_names(engine, mes.from_user.id):
return await send_tmp_message(bot, mes.chat.id, "Нет такого аккаунта") return await send_tmp_message(bot, mes.chat.id, "Нет такого аккаунта")
bot_mes = await bot.send_message(mes.chat.id, "Отправьте мастер пароль") bot_mes = await bot.send_message(mes.chat.id, "Отправьте мастер пароль")
@ -378,7 +376,7 @@ async def _get_account3(
if text == "/cancel": if text == "/cancel":
return await send_tmp_message(bot, mes.chat.id, "Успешная отмена") return await send_tmp_message(bot, mes.chat.id, "Успешная отмена")
master_password = database.get.get_master_pass( master_password = db.get.get_master_pass(
engine, engine,
mes.from_user.id, mes.from_user.id,
) )
@ -390,17 +388,14 @@ async def _get_account3(
"Не подходит мастер пароль", "Не подходит мастер пароль",
) )
account = database.get.get_account_info(engine, mes.from_user.id, name) account = db.get.get_account_info(engine, mes.from_user.id, name)
account = encryption.other_accounts.decrypt( account = encryption.accounts.decrypt(account, text)
account,
text,
)
await send_deleteable_message( await send_deleteable_message(
bot, bot,
mes.chat.id, mes.chat.id,
f"Название:\n{escape(account.name)}\n" f"Название:\n`{account.name}`\n"
f"Логин:\n`{account.login}`\nПароль:\n`{account.password}`\nНажмите " f"Логин:\n`{account.login}`\nПароль:\n`{account.password}`\nНажмите "
"на логин или пароль, чтобы скопировать", "на название, логин или пароль, чтобы скопировать",
) )
del text, mes del text, mes
@ -414,11 +409,11 @@ async def delete_account(
) -> None: ) -> None:
await base_handler(bot, mes) await base_handler(bot, mes)
master_pass = database.get.get_master_pass(engine, mes.from_user.id) master_pass = db.get.get_master_pass(engine, mes.from_user.id)
if master_pass is None: if master_pass is None:
return await send_tmp_message(bot, mes.chat.id, "Нет мастер пароля") return await send_tmp_message(bot, mes.chat.id, "Нет мастер пароля")
accounts = database.get.get_accounts( accounts = db.get.get_account_names(
engine, engine,
mes.from_user.id, mes.from_user.id,
to_sort=True, to_sort=True,
@ -447,7 +442,7 @@ async def _delete_account2(
if text == "/cancel": if text == "/cancel":
return await send_tmp_message(bot, mes.chat.id, "Успешная отмена") return await send_tmp_message(bot, mes.chat.id, "Успешная отмена")
if text not in database.get.get_accounts(engine, mes.from_user.id): if text not in db.get.get_account_names(engine, mes.from_user.id):
return await send_tmp_message(bot, mes.chat.id, "Нет такого аккаунта") return await send_tmp_message(bot, mes.chat.id, "Нет такого аккаунта")
bot_mes = await bot.send_message( bot_mes = await bot.send_message(
@ -474,7 +469,7 @@ async def _delete_account3(
if text != "YES": if text != "YES":
return await send_tmp_message(bot, mes.chat.id, "Успешная отмена") return await send_tmp_message(bot, mes.chat.id, "Успешная отмена")
database.delete.delete_account(engine, mes.from_user.id, account_name) db.delete.delete_account(engine, mes.from_user.id, account_name)
await send_tmp_message(bot, mes.chat.id, "Аккаунт удалён") await send_tmp_message(bot, mes.chat.id, "Аккаунт удалён")
@ -498,7 +493,7 @@ async def help_command(bot: AsyncTeleBot, mes: Message) -> None:
async def export(bot: AsyncTeleBot, engine: Engine, mes: Message) -> None: async def export(bot: AsyncTeleBot, engine: Engine, mes: Message) -> None:
await base_handler(bot, mes) await base_handler(bot, mes)
master_password_from_db = database.get.get_master_pass( master_password_from_db = db.get.get_master_pass(
engine, engine,
mes.from_user.id, mes.from_user.id,
) )
@ -506,7 +501,7 @@ async def export(bot: AsyncTeleBot, engine: Engine, mes: Message) -> None:
if master_password_from_db is None: if master_password_from_db is None:
return await send_tmp_message(bot, mes.chat.id, "Нет мастер пароля") return await send_tmp_message(bot, mes.chat.id, "Нет мастер пароля")
if not database.get.get_accounts(engine, mes.from_user.id): if not db.get.get_account_names(engine, mes.from_user.id):
return await send_tmp_message(bot, mes.chat.id, "Нет аккаунтов") return await send_tmp_message(bot, mes.chat.id, "Нет аккаунтов")
bot_mes = await bot.send_message(mes.chat.id, "Отправьте мастер пароль") bot_mes = await bot.send_message(mes.chat.id, "Отправьте мастер пароль")
@ -522,7 +517,7 @@ async def _export2(
if text == "/cancel": if text == "/cancel":
return await send_tmp_message(bot, mes.chat.id, "Успешная отмена") return await send_tmp_message(bot, mes.chat.id, "Успешная отмена")
master_password = database.get.get_master_pass( master_password = db.get.get_master_pass(
engine, engine,
mes.from_user.id, mes.from_user.id,
) )
@ -533,8 +528,19 @@ async def _export2(
"Не подходит мастер пароль", "Не подходит мастер пароль",
) )
accounts = database.get.get_all_accounts(engine, mes.from_user.id) accounts = db.get.get_accounts(engine, mes.from_user.id)
accounts = encryption.other_accounts.decrypt_multiple(accounts, text) with ProcessPoolExecutor() as pool:
loop = asyncio.get_running_loop()
tasks = []
for account in accounts:
function = functools.partial(
encryption.accounts.decrypt,
account,
text,
)
tasks.append(loop.run_in_executor(pool, function))
accounts = await asyncio.gather(*tasks)
accounts.sort(key=lambda account: account.name)
json_io = accounts_to_json(accounts) json_io = accounts_to_json(accounts)
await bot.send_document( await bot.send_document(
mes.chat.id, mes.chat.id,
@ -552,7 +558,7 @@ async def import_accounts(
mes: Message, mes: Message,
) -> None: ) -> None:
await base_handler(bot, mes) await base_handler(bot, mes)
master_password_from_db = database.get.get_master_pass( master_password_from_db = db.get.get_master_pass(
engine, engine,
mes.from_user.id, mes.from_user.id,
) )
@ -620,7 +626,7 @@ async def _import3(
if text == "/cancel": if text == "/cancel":
return await send_tmp_message(bot, mes.chat.id, "Успешная отмена") return await send_tmp_message(bot, mes.chat.id, "Успешная отмена")
master_password = database.get.get_master_pass( master_password = db.get.get_master_pass(
engine, engine,
mes.from_user.id, mes.from_user.id,
) )
@ -638,11 +644,8 @@ async def _import3(
if not check_account(account): if not check_account(account):
failed.append(account.name) failed.append(account.name)
continue continue
account = encryption.other_accounts.encrypt( account = encryption.accounts.encrypt(account, text)
account, result = db.add.add_account(engine, account)
text,
)
result = database.add.add_account(engine, account)
if not result: if not result:
failed.append(account.name) failed.append(account.name)

View File

@ -17,7 +17,7 @@ def get_master_pass(
return result return result
def get_accounts( def get_account_names(
engine: Engine, engine: Engine,
user_id: int, user_id: int,
*, *,
@ -34,14 +34,10 @@ def get_accounts(
return result return result
def get_all_accounts(engine: Engine, user_id: int) -> list[models.Account]: def get_accounts(engine: Engine, user_id: int) -> list[models.Account]:
"""Returns a list of accounts of a user""" """Returns a list of accounts of a user"""
statement = ( statement = sqlmodel.select(models.Account).where(
sqlmodel.select(models.Account) models.Account.user_id == user_id,
.where(
models.Account.user_id == user_id,
)
.order_by(models.Account.name)
) )
with sqlmodel.Session(engine) as session: with sqlmodel.Session(engine) as session:
result = session.exec(statement).fetchall() result = session.exec(statement).fetchall()

View File

@ -1,3 +1,3 @@
from . import master_pass, other_accounts from . import accounts, master_pass
__all__ = ["master_pass", "other_accounts"] __all__ = ["master_pass", "accounts"]

View File

@ -1,13 +1,12 @@
import base64 import base64
import os import os
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 ..database.models import Account from ..db.models import Account
from ..decrypted_account import DecryptedAccount from ..decrypted_account import DecryptedAccount
@ -71,12 +70,3 @@ def decrypt(
login=login, login=login,
password=password, password=password,
) )
def decrypt_multiple(
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:
yield decrypt(account, master_pass)

View File

@ -3,7 +3,7 @@ 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 from ..db.models import MasterPass
MEMORY_USAGE = 2**14 MEMORY_USAGE = 2**14