Compare commits

..

No commits in common. "main" and "1.3.1" have entirely different histories.
main ... 1.3.1

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:
result = _Accounts(
accounts=[_Account.from_usual_account(i) for i in accounts],
).json(ensure_ascii=False, indent=2)
).json(ensure_ascii=False)
return result

View File

@ -3,13 +3,13 @@ import functools
from sqlalchemy.future import Engine
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:
bot = AsyncTeleBot(token, exception_handler=exception_handler.Handler)
bot = AsyncTeleBot(token)
bot.register_message_handler(
functools.partial(message_handlers.set_master_password, bot, engine),
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 functools
import gc
import itertools
from concurrent.futures import ProcessPoolExecutor
import telebot
@ -58,31 +57,21 @@ async def get_accounts(
async def delete_all(bot: AsyncTeleBot, engine: Engine, mes: Message) -> None:
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(
mes.chat.id,
"Вы действительно хотите удалить все ваши аккаунты? Это действие "
"нельзя отменить. "
"Отправьте мастер пароль для подтверждения",
)
register_state(
mes, functools.partial(_delete_all2, bot, engine, master_pass, bot_mes)
"Отправьте YES для подтверждения",
)
register_state(mes, functools.partial(_delete_all2, bot, engine, bot_mes))
async def _delete_all2(
bot: AsyncTeleBot,
engine: Engine,
master_pass: db.models.MasterPass,
prev_mes: Message,
mes: Message,
bot: AsyncTeleBot, engine: Engine, prev_mes: Message, mes: Message
) -> None:
await base_handler(bot, mes, prev_mes)
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.delete_master_pass(engine, mes.from_user.id)
await send_tmp_message(
@ -95,7 +84,7 @@ async def _delete_all2(
await send_tmp_message(
bot,
mes.chat.id,
"Вы отправили не верный мастер пароль, ничего не удалено",
"Вы отправили не YES, ничего не удалено",
)
@ -144,8 +133,7 @@ async def reset_master_pass(
) -> None:
await base_handler(bot, mes)
master_pass = db.get.get_master_pass(engine, mes.from_user.id)
if master_pass is None:
if db.get.get_master_pass(engine, mes.from_user.id) is None:
return await send_tmp_message(
bot,
mes.chat.id,
@ -154,48 +142,17 @@ async def reset_master_pass(
bot_mes = await bot.send_message(
mes.chat.id,
"Отправьте текущий мастер пароль",
"Отправьте новый мастер пароль, осторожно, все текущие аккаунты "
"будут удалены навсегда",
)
register_state(
mes,
functools.partial(
_reset_master_pass2,
bot,
engine,
master_pass,
bot_mes,
),
functools.partial(_reset_master_pass2, bot, engine, bot_mes),
)
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
) -> None:
await base_handler(bot, mes, prev_mes)
@ -470,14 +427,13 @@ async def delete_account(
register_state(
mes,
functools.partial(_delete_account2, bot, engine, master_pass, bot_mes),
functools.partial(_delete_account2, bot, engine, bot_mes),
)
async def _delete_account2(
bot: AsyncTeleBot,
engine: Engine,
master_pass: db.models.MasterPass,
prev_mes: Message,
mes: Message,
):
@ -491,36 +447,27 @@ async def _delete_account2(
bot_mes = await bot.send_message(
mes.from_user.id,
f'Вы уверены, что хотите удалить аккаунт "{text}"?\nОтправьте мастер '
"пароль для подтверждения",
f'Вы уверены, что хотите удалить аккаунт "{text}"?\nОтправьте YES для '
"подтверждения",
)
register_state(
mes,
functools.partial(
_delete_account3,
bot,
engine,
master_pass,
bot_mes,
text,
),
functools.partial(_delete_account3, bot, engine, bot_mes, text),
)
async def _delete_account3(
bot: AsyncTeleBot,
engine: Engine,
master_pass: db.models.MasterPass,
prev_mes: Message,
account_name: str,
mes: Message,
) -> None:
await base_handler(bot, mes, prev_mes)
text = mes.text.strip()
if not encryption.master_pass.check_master_pass(text, master_pass):
await send_tmp_message(bot, mes.chat.id, "Неверный пароль")
return
if text != "YES":
return await send_tmp_message(bot, mes.chat.id, "Успешная отмена")
db.delete.delete_account(engine, mes.from_user.id, account_name)
await send_tmp_message(bot, mes.chat.id, "Аккаунт удалён")
@ -593,6 +540,7 @@ async def _export2(
)
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)
await bot.send_document(
mes.chat.id,
@ -692,26 +640,14 @@ async def _import3(
# List of names of accounts, which failed to be added to the database
# or failed the tests
failed: list[str] = []
tasks: list[asyncio.Future[db.models.Account]] = []
loop = asyncio.get_running_loop()
with ProcessPoolExecutor() as pool:
for account in accounts:
if not check_account(account):
failed.append(account.name)
continue
function = functools.partial(
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))
account = encryption.accounts.encrypt(account, text)
result = db.add.add_account(engine, account)
if not result:
failed.append(account.name)
if failed:
await send_deleteable_message(
@ -720,7 +656,7 @@ async def _import3(
else:
await send_tmp_message(bot, mes.chat.id, "Успех")
del text, mes, accounts, function, tasks, failed_accounts
del text, mes, accounts
gc.collect()
@ -742,20 +678,10 @@ async def message_handler(bot: AsyncTeleBot, mes: Message) -> None:
await delete_message(bot, mes)
if mes.text.strip() == "/cancel":
await send_tmp_message(bot, mes.chat.id, "Нет активного действия")
return
await send_tmp_message(
bot,
mes.chat.id,
"Вы отправили не корректное сообщение",
)
return
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
def _add_model(
session: sqlmodel.Session, model: models.Account | models.MasterPass
) -> bool:
"""Adds model to the session. Returns true on success,
def add_account(engine: Engine, account: models.Account) -> bool:
"""Adds account to the database. Returns true on success,
false otherwise"""
try:
session.add(model)
with sqlmodel.Session(engine) as session:
session.add(account)
session.commit()
except IntegrityError:
return False
else:
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:
"""Adds master password the database. Returns true on success,
false otherwise"""
try:
with sqlmodel.Session(engine) as session:
result = _add_model(session, master_pass)
session.add(master_pass)
session.commit()
return result
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
except IntegrityError:
return False
else:
return True

View File

@ -36,13 +36,9 @@ def get_account_names(
def get_accounts(engine: Engine, user_id: int) -> list[models.Account]:
"""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,
)
.order_by(models.Account.name)
)
with sqlmodel.Session(engine) as session:
result = session.exec(statement).fetchall()
return result

View File

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

View File

@ -1,43 +1,26 @@
import base64
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.ciphers.aead import ChaCha20Poly1305
from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
from ..db.models import Account
from ..decrypted_account import DecryptedAccount
class Cipher:
def __init__(self, key: bytes) -> None:
self._chacha = ChaCha20Poly1305(key)
@classmethod
def generate_cipher(cls, salt: bytes, password: bytes) -> Self:
"""Generates cipher which uses key derived from a given password"""
def _generate_key(salt: bytes, master_pass: bytes) -> bytes:
"""Generates key for fernet encryption"""
kdf = PBKDF2HMAC(
algorithm=hashes.SHA256(),
length=32,
salt=salt,
iterations=100000,
backend=default_backend(),
)
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,
)
key = base64.urlsafe_b64encode(kdf.derive(master_pass))
return key
def encrypt(
@ -46,10 +29,15 @@ def encrypt(
) -> Account:
"""Encrypts account using master password and returns Account object"""
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_password = cipher.encrypt(account.password.encode("utf-8"))
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,
@ -66,10 +54,15 @@ def decrypt(
) -> DecryptedAccount:
"""Decrypts account using master password and returns
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")
password = cipher.decrypt(account.enc_password).decode("utf-8")
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,