Compare commits

..

No commits in common. "972c5577f4aafd820308cd79024b06b52f93ba38" and "c7675c231fbacfb7c4e308d01d4a5878ce7b2e38" have entirely different histories.

8 changed files with 82 additions and 202 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, indent=2) ).json(ensure_ascii=False)
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, exception_handler, message_handlers from . import callback_handlers, message_handlers
__all__ = ["callback_handlers", "exception_handler", "message_handlers"] __all__ = ["callback_handlers", "message_handlers"]
def create_bot(token: str, engine: Engine) -> AsyncTeleBot: def create_bot(token: str, engine: Engine) -> AsyncTeleBot:
bot = AsyncTeleBot(token, exception_handler=exception_handler.Handler) bot = AsyncTeleBot(token)
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

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

View File

@ -1,7 +1,6 @@
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
@ -58,31 +57,21 @@ async def get_accounts(
async def delete_all(bot: AsyncTeleBot, engine: Engine, mes: Message) -> None: async def delete_all(bot: AsyncTeleBot, engine: Engine, mes: Message) -> None:
await base_handler(bot, mes) await base_handler(bot, mes)
master_pass = db.get.get_master_pass(engine, mes.from_user.id)
if master_pass is None:
await send_tmp_message(bot, mes.chat.id, "У вас нет мастер пароля")
return
bot_mes = await bot.send_message( bot_mes = await bot.send_message(
mes.chat.id, mes.chat.id,
"Вы действительно хотите удалить все ваши аккаунты? Это действие " "Вы действительно хотите удалить все ваши аккаунты? Это действие "
"нельзя отменить. " "нельзя отменить. "
"Отправьте мастер пароль для подтверждения", "Отправьте YES для подтверждения",
)
register_state(
mes, functools.partial(_delete_all2, bot, engine, master_pass, bot_mes)
) )
register_state(mes, functools.partial(_delete_all2, bot, engine, bot_mes))
async def _delete_all2( async def _delete_all2(
bot: AsyncTeleBot, bot: AsyncTeleBot, engine: Engine, prev_mes: Message, mes: Message
engine: Engine,
master_pass: db.models.MasterPass,
prev_mes: Message,
mes: Message,
) -> None: ) -> None:
await base_handler(bot, mes, prev_mes) await base_handler(bot, mes, prev_mes)
text = mes.text.strip() text = mes.text.strip()
if encryption.master_pass.check_master_pass(text, master_pass): if text == "YES":
db.delete.purge_accounts(engine, mes.from_user.id) db.delete.purge_accounts(engine, mes.from_user.id)
db.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(
@ -95,7 +84,7 @@ async def _delete_all2(
await send_tmp_message( await send_tmp_message(
bot, bot,
mes.chat.id, mes.chat.id,
"Вы отправили не верный мастер пароль, ничего не удалено", "Вы отправили не YES, ничего не удалено",
) )
@ -144,8 +133,7 @@ async def reset_master_pass(
) -> None: ) -> None:
await base_handler(bot, mes) await base_handler(bot, mes)
master_pass = db.get.get_master_pass(engine, mes.from_user.id) if db.get.get_master_pass(engine, mes.from_user.id) is None:
if master_pass is None:
return await send_tmp_message( return await send_tmp_message(
bot, bot,
mes.chat.id, mes.chat.id,
@ -154,48 +142,17 @@ async def reset_master_pass(
bot_mes = await bot.send_message( bot_mes = await bot.send_message(
mes.chat.id, mes.chat.id,
"Отправьте текущий мастер пароль", "Отправьте новый мастер пароль, осторожно, все текущие аккаунты "
"будут удалены навсегда",
) )
register_state( register_state(
mes, mes,
functools.partial( functools.partial(_reset_master_pass2, bot, engine, bot_mes),
_reset_master_pass2,
bot,
engine,
master_pass,
bot_mes,
),
) )
async def _reset_master_pass2( async def _reset_master_pass2(
bot: AsyncTeleBot,
engine: Engine,
master_pass: db.models.MasterPass,
prev_mes: Message,
mes: Message,
) -> None:
await base_handler(bot, mes, prev_mes)
text = mes.text.strip()
if text == "/cancel":
await send_tmp_message(bot, mes.chat.id, "Успешная отмена")
if not encryption.master_pass.check_master_pass(text, master_pass):
await send_tmp_message(bot, mes.chat.id, "Неверный мастер пароль")
return
bot_mes = await bot.send_message(
mes.chat.id,
"Отправьте новый мастер пароль. Осторожно, все аккаунты будут удалены",
)
register_state(
mes,
functools.partial(_reset_master_pass3, bot, engine, bot_mes),
)
async def _reset_master_pass3(
bot: AsyncTeleBot, engine: Engine, prev_mes: Message, mes: Message bot: AsyncTeleBot, engine: Engine, prev_mes: Message, mes: Message
) -> None: ) -> None:
await base_handler(bot, mes, prev_mes) await base_handler(bot, mes, prev_mes)
@ -470,14 +427,13 @@ async def delete_account(
register_state( register_state(
mes, mes,
functools.partial(_delete_account2, bot, engine, master_pass, bot_mes), functools.partial(_delete_account2, bot, engine, bot_mes),
) )
async def _delete_account2( async def _delete_account2(
bot: AsyncTeleBot, bot: AsyncTeleBot,
engine: Engine, engine: Engine,
master_pass: db.models.MasterPass,
prev_mes: Message, prev_mes: Message,
mes: Message, mes: Message,
): ):
@ -491,36 +447,27 @@ async def _delete_account2(
bot_mes = await bot.send_message( bot_mes = await bot.send_message(
mes.from_user.id, mes.from_user.id,
f'Вы уверены, что хотите удалить аккаунт "{text}"?\nОтправьте мастер ' f'Вы уверены, что хотите удалить аккаунт "{text}"?\nОтправьте YES для '
"пароль для подтверждения", "подтверждения",
) )
register_state( register_state(
mes, mes,
functools.partial( functools.partial(_delete_account3, bot, engine, bot_mes, text),
_delete_account3,
bot,
engine,
master_pass,
bot_mes,
text,
),
) )
async def _delete_account3( async def _delete_account3(
bot: AsyncTeleBot, bot: AsyncTeleBot,
engine: Engine, engine: Engine,
master_pass: db.models.MasterPass,
prev_mes: Message, prev_mes: Message,
account_name: str, account_name: str,
mes: Message, mes: Message,
) -> None: ) -> None:
await base_handler(bot, mes, prev_mes) await base_handler(bot, mes, prev_mes)
text = mes.text.strip() text = mes.text.strip()
if not encryption.master_pass.check_master_pass(text, master_pass): if text != "YES":
await send_tmp_message(bot, mes.chat.id, "Неверный пароль") return await send_tmp_message(bot, mes.chat.id, "Успешная отмена")
return
db.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, "Аккаунт удалён")
@ -593,6 +540,7 @@ 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,
@ -692,26 +640,14 @@ 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]] = [] for account in accounts:
loop = asyncio.get_running_loop() if not check_account(account):
with ProcessPoolExecutor() as pool: failed.append(account.name)
for account in accounts: continue
if not check_account(account): account = encryption.accounts.encrypt(account, text)
failed.append(account.name) result = db.add.add_account(engine, account)
continue if not result:
function = functools.partial( failed.append(account.name)
encryption.accounts.encrypt,
account,
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(
@ -720,7 +656,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, function, tasks, failed_accounts del text, mes, accounts
gc.collect() gc.collect()
@ -742,20 +678,10 @@ 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,
"Вы отправили не корректное сообщение", "Вы отправили не корректное сообщение",
) )
return return
await handler(mes)
try:
await handler(mes)
except Exception:
await send_tmp_message(
bot,
mes.chat.id,
"Произошла непредвиденная ошибка",
)
raise

View File

@ -5,42 +5,27 @@ from sqlalchemy.future import Engine
from . import models from . import models
def _add_model( def add_account(engine: Engine, account: models.Account) -> bool:
session: sqlmodel.Session, model: models.Account | models.MasterPass """Adds account to the database. Returns true on success,
) -> bool:
"""Adds model to the session. Returns true on success,
false otherwise""" false otherwise"""
try: try:
session.add(model) with sqlmodel.Session(engine) as session:
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"""
with sqlmodel.Session(engine) as session: try:
result = _add_model(session, master_pass) with sqlmodel.Session(engine) as session:
session.commit() session.add(master_pass)
return result session.commit()
except IntegrityError:
return False
def add_accounts( else:
engine: Engine, return True
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,12 +36,8 @@ 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 = ( 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

@ -4,21 +4,13 @@ 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( sa_column=sqlmodel.Column(sqlmodel.INT(), primary_key=True)
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)
max_length=64,
min_length=64,
) )
password_hash: bytes = sqlmodel.Field( password_hash: bytes = sqlmodel.Field(
sa_column=sqlmodel.Column(sqlmodel.BINARY(128), nullable=False), sa_column=sqlmodel.Column(sqlmodel.BINARY(128), nullable=False)
max_length=128,
min_length=128,
) )
@ -26,17 +18,13 @@ class Account(sqlmodel.SQLModel, table=True):
__tablename__ = "accounts" __tablename__ = "accounts"
__table_args__ = (sqlmodel.PrimaryKeyConstraint("user_id", "name"),) __table_args__ = (sqlmodel.PrimaryKeyConstraint("user_id", "name"),)
user_id: int = sqlmodel.Field() user_id: int = sqlmodel.Field()
name: str = sqlmodel.Field(max_length=256) name: str = sqlmodel.Field(max_length=255)
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)
max_length=64,
min_length=64,
) )
enc_login: bytes = sqlmodel.Field( enc_login: bytes = sqlmodel.Field(
sa_column=sqlmodel.Column(sqlmodel.VARBINARY(256), nullable=False), sa_column=sqlmodel.Column(sqlmodel.VARBINARY(256), nullable=False)
max_length=256,
) )
enc_password: bytes = sqlmodel.Field( enc_password: bytes = sqlmodel.Field(
sa_column=sqlmodel.Column(sqlmodel.VARBINARY(256), nullable=False), sa_column=sqlmodel.Column(sqlmodel.VARBINARY(256), nullable=False)
max_length=256,
) )

View File

@ -1,43 +1,26 @@
import base64
import os import os
from typing import Self
from cryptography.fernet import Fernet
from cryptography.hazmat.backends import default_backend
from cryptography.hazmat.primitives import hashes from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.ciphers.aead import ChaCha20Poly1305
from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
from ..db.models import Account from ..db.models import Account
from ..decrypted_account import DecryptedAccount from ..decrypted_account import DecryptedAccount
class Cipher: def _generate_key(salt: bytes, master_pass: bytes) -> bytes:
def __init__(self, key: bytes) -> None: """Generates key for fernet encryption"""
self._chacha = ChaCha20Poly1305(key) kdf = PBKDF2HMAC(
algorithm=hashes.SHA256(),
@classmethod length=32,
def generate_cipher(cls, salt: bytes, password: bytes) -> Self: salt=salt,
"""Generates cipher which uses key derived from a given password""" iterations=100000,
kdf = PBKDF2HMAC( backend=default_backend(),
algorithm=hashes.SHA256(), )
length=32, key = base64.urlsafe_b64encode(kdf.derive(master_pass))
salt=salt, return key
iterations=100000,
)
return cls(kdf.derive(password))
def encrypt(self, data: bytes) -> bytes:
nonce = os.urandom(12)
return nonce + self._chacha.encrypt(
nonce,
data,
associated_data=None,
)
def decrypt(self, data: bytes) -> bytes:
return self._chacha.decrypt(
nonce=data[:12],
data=data[12:],
associated_data=None,
)
def encrypt( def encrypt(
@ -46,10 +29,15 @@ def encrypt(
) -> Account: ) -> Account:
"""Encrypts account using master password and returns Account object""" """Encrypts account using master password and returns Account object"""
salt = os.urandom(64) salt = os.urandom(64)
cipher = Cipher.generate_cipher(salt, master_pass.encode("utf-8")) key = _generate_key(salt, master_pass.encode("utf-8"))
f = Fernet(key)
enc_login = cipher.encrypt(account.login.encode("utf-8")) enc_login = base64.urlsafe_b64decode(
enc_password = cipher.encrypt(account.password.encode("utf-8")) f.encrypt(account.login.encode("utf-8")),
)
enc_password = base64.urlsafe_b64decode(
f.encrypt(account.password.encode("utf-8")),
)
return Account( return Account(
user_id=account.user_id, user_id=account.user_id,
@ -66,10 +54,15 @@ def decrypt(
) -> DecryptedAccount: ) -> DecryptedAccount:
"""Decrypts account using master password and returns """Decrypts account using master password and returns
DecryptedAccount object""" DecryptedAccount object"""
cipher = Cipher.generate_cipher(account.salt, master_pass.encode("utf-8")) key = _generate_key(account.salt, master_pass.encode("utf-8"))
f = Fernet(key)
login = cipher.decrypt(account.enc_login).decode("utf-8") login = f.decrypt(
password = cipher.decrypt(account.enc_password).decode("utf-8") base64.urlsafe_b64encode(account.enc_login),
).decode("utf-8")
password = f.decrypt(
base64.urlsafe_b64encode(account.enc_password),
).decode("utf-8")
return DecryptedAccount( return DecryptedAccount(
user_id=account.user_id, user_id=account.user_id,