7 Commits

Author SHA1 Message Date
d79b57b1f0 Added return statement after sending a message about that there is no active message 2023-01-05 14:52:45 +03:00
157c2c4aa2 _import3 no longer blocks event loop during decryption
Now running encrytion of accounts in ProcessPoolExecutor
2023-01-05 14:11:29 +03:00
f4a5f51b23 Made more verbose exception handler for the bot for easier debugging 2023-01-05 13:47:31 +03:00
6bc8eb1413 db.add changes
Added _add_model helper function to reduce code duplication
Added add_accounts for future use
2023-01-05 13:19:01 +03:00
9f64305050 Moved sorting back to the get_accounts 2023-01-05 13:03:44 +03:00
4954f39a91 Added indentation into exported json files 2023-01-05 13:02:50 +03:00
3edeb86b6c Disabled autoincrement in master_passwords table 2023-01-05 13:01:28 +03:00
7 changed files with 74 additions and 30 deletions

View File

@ -35,7 +35,7 @@ class _Accounts(pydantic.BaseModel):
def _accounts_list_to_json(accounts: Iterable[DecryptedAccount]) -> str: def _accounts_list_to_json(accounts: Iterable[DecryptedAccount]) -> str:
result = _Accounts( result = _Accounts(
accounts=[_Account.from_usual_account(i) for i in accounts], accounts=[_Account.from_usual_account(i) for i in accounts],
).json(ensure_ascii=False) ).json(ensure_ascii=False, indent=2)
return result return result

View File

@ -3,13 +3,13 @@ import functools
from sqlalchemy.future import Engine from sqlalchemy.future import Engine
from telebot.async_telebot import AsyncTeleBot from telebot.async_telebot import AsyncTeleBot
from . import callback_handlers, message_handlers from . import callback_handlers, exception_handler, message_handlers
__all__ = ["callback_handlers", "message_handlers"] __all__ = ["callback_handlers", "exception_handler", "message_handlers"]
def create_bot(token: str, engine: Engine) -> AsyncTeleBot: def create_bot(token: str, engine: Engine) -> AsyncTeleBot:
bot = AsyncTeleBot(token) bot = AsyncTeleBot(token, exception_handler=exception_handler.Handler)
bot.register_message_handler( bot.register_message_handler(
functools.partial(message_handlers.set_master_password, bot, engine), functools.partial(message_handlers.set_master_password, bot, engine),
commands=["set_master_pass"], commands=["set_master_pass"],

View File

@ -0,0 +1,8 @@
import traceback
from typing import Type
class Handler:
@staticmethod
def handle(exc: Type[BaseException]) -> None:
traceback.print_exception(exc)

View File

@ -1,6 +1,7 @@
import asyncio import asyncio
import functools import functools
import gc import gc
import itertools
from concurrent.futures import ProcessPoolExecutor from concurrent.futures import ProcessPoolExecutor
import telebot import telebot
@ -540,7 +541,6 @@ async def _export2(
) )
tasks.append(loop.run_in_executor(pool, function)) tasks.append(loop.run_in_executor(pool, function))
accounts = await asyncio.gather(*tasks) 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,
@ -640,14 +640,26 @@ async def _import3(
# List of names of accounts, which failed to be added to the database # List of names of accounts, which failed to be added to the database
# or failed the tests # or failed the tests
failed: list[str] = [] failed: list[str] = []
tasks: list[asyncio.Future[db.models.Account]] = []
loop = asyncio.get_running_loop()
with ProcessPoolExecutor() as pool:
for account in accounts: for account in accounts:
if not check_account(account): if not check_account(account):
failed.append(account.name) failed.append(account.name)
continue continue
account = encryption.accounts.encrypt(account, text) function = functools.partial(
result = db.add.add_account(engine, account) encryption.accounts.encrypt,
if not result: account,
failed.append(account.name) text,
)
tasks.append(loop.run_in_executor(pool, function))
enc_accounts: list[db.models.Account] = await asyncio.gather(*tasks)
results = db.add.add_accounts(engine, enc_accounts)
failed_accounts = itertools.compress(
enc_accounts, (not result for result in results)
)
failed.extend((account.name for account in failed_accounts))
if failed: if failed:
await send_deleteable_message( await send_deleteable_message(
@ -656,7 +668,7 @@ async def _import3(
else: else:
await send_tmp_message(bot, mes.chat.id, "Успех") await send_tmp_message(bot, mes.chat.id, "Успех")
del text, mes, accounts del text, mes, accounts, function, tasks, failed_accounts
gc.collect() gc.collect()
@ -678,6 +690,7 @@ async def message_handler(bot: AsyncTeleBot, mes: Message) -> None:
await delete_message(bot, mes) await delete_message(bot, mes)
if mes.text.strip() == "/cancel": if mes.text.strip() == "/cancel":
await send_tmp_message(bot, mes.chat.id, "Нет активного действия") await send_tmp_message(bot, mes.chat.id, "Нет активного действия")
return
await send_tmp_message( await send_tmp_message(
bot, bot,
mes.chat.id, mes.chat.id,

View File

@ -5,27 +5,42 @@ from sqlalchemy.future import Engine
from . import models from . import models
def add_account(engine: Engine, account: models.Account) -> bool: def _add_model(
"""Adds account to the database. Returns true on success, session: sqlmodel.Session, model: models.Account | models.MasterPass
) -> bool:
"""Adds model to the session. Returns true on success,
false otherwise""" false otherwise"""
try: try:
with sqlmodel.Session(engine) as session: session.add(model)
session.add(account)
session.commit()
except IntegrityError: except IntegrityError:
return False return False
else: else:
return True return True
def add_account(engine: Engine, account: models.Account) -> bool:
"""Adds account to the database. Returns true on success,
false otherwise"""
with sqlmodel.Session(engine) as session:
result = _add_model(session, account)
session.commit()
return result
def add_master_pass(engine: Engine, master_pass: models.MasterPass) -> bool: def add_master_pass(engine: Engine, master_pass: models.MasterPass) -> bool:
"""Adds master password the database. Returns true on success, """Adds master password the database. Returns true on success,
false otherwise""" false otherwise"""
try:
with sqlmodel.Session(engine) as session: with sqlmodel.Session(engine) as session:
session.add(master_pass) result = _add_model(session, master_pass)
session.commit() session.commit()
except IntegrityError: return result
return False
else:
return True def add_accounts(
engine: Engine,
accounts: list[models.Account],
) -> list[bool]:
with sqlmodel.Session(engine) as session:
result = [_add_model(session, account) for account in accounts]
session.commit()
return result

View File

@ -36,9 +36,13 @@ def get_account_names(
def get_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 = sqlmodel.select(models.Account).where( statement = (
sqlmodel.select(models.Account)
.where(
models.Account.user_id == user_id, 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()
return result return result

View File

@ -4,7 +4,11 @@ import sqlmodel
class MasterPass(sqlmodel.SQLModel, table=True): class MasterPass(sqlmodel.SQLModel, table=True):
__tablename__ = "master_passwords" __tablename__ = "master_passwords"
user_id: int = sqlmodel.Field( user_id: int = sqlmodel.Field(
sa_column=sqlmodel.Column(sqlmodel.INT(), primary_key=True) sa_column=sqlmodel.Column(
sqlmodel.INT(),
primary_key=True,
autoincrement=False,
)
) )
salt: bytes = sqlmodel.Field( salt: bytes = sqlmodel.Field(
sa_column=sqlmodel.Column(sqlmodel.BINARY(64), nullable=False) sa_column=sqlmodel.Column(sqlmodel.BINARY(64), nullable=False)