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
This commit is contained in:
StNicolay 2022-12-25 18:01:12 +03:00
parent bbc9650357
commit 3f744723a9
12 changed files with 195 additions and 186 deletions

View File

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

View File

@ -6,8 +6,9 @@ from . import (
account_checks, account_checks,
account_parsing, account_parsing,
bot, bot,
encryption, classes,
database, database,
encryption,
generate_password, generate_password,
) )
@ -15,6 +16,7 @@ __all__ = [
"account_checks", "account_checks",
"account_parsing", "account_parsing",
"bot", "bot",
"classes",
"encryption", "encryption",
"database", "database",
"generate_password", "generate_password",

View File

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

View File

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

View File

@ -5,7 +5,7 @@ import time
import telebot import telebot
from sqlalchemy.future import Engine from sqlalchemy.future import Engine
from .. import encryption, database, generate_password from .. import database, encryption, generate_password
from ..account_checks import ( from ..account_checks import (
check_account, check_account,
check_account_name, check_account_name,
@ -13,6 +13,7 @@ from ..account_checks import (
check_password, check_password,
) )
from ..account_parsing import accounts_to_json, json_to_accounts from ..account_parsing import accounts_to_json, json_to_accounts
from ..classes import DecryptedAccount
Message = telebot.types.Message Message = telebot.types.Message
@ -31,9 +32,9 @@ def _send_tmp_message(
def _base_handler( def _base_handler(
bot: telebot.TeleBot, mes: Message, prev_mes: Message | None = None bot: telebot.TeleBot, mes: Message, prev_mes: Message | None = None
) -> None: ) -> None:
bot.delete_message(mes.chat.id, mes.id) bot.delete_message(mes.chat.id, mes.id, timeout=5)
if prev_mes is not None: if prev_mes is not None:
bot.delete_message(prev_mes.chat.id, prev_mes.id) bot.delete_message(prev_mes.chat.id, prev_mes.id, timeout=5)
def get_accounts(bot: telebot.TeleBot, engine: Engine, mes: Message) -> None: def get_accounts(bot: telebot.TeleBot, engine: Engine, mes: Message) -> None:
@ -68,11 +69,11 @@ def delete_all(bot: telebot.TeleBot, engine: Engine, mes: Message) -> None:
"Отправьте YES для подтверждения", "Отправьте YES для подтверждения",
) )
bot.register_next_step_handler( bot.register_next_step_handler(
mes, functools.partial(_delete_all, bot, engine, bot_mes) mes, functools.partial(_delete_all2, bot, engine, bot_mes)
) )
def _delete_all( def _delete_all2(
bot: telebot.TeleBot, engine: Engine, prev_mes: Message, mes: Message bot: telebot.TeleBot, engine: Engine, prev_mes: Message, mes: Message
) -> None: ) -> None:
_base_handler(bot, mes, prev_mes) _base_handler(bot, mes, prev_mes)
@ -115,15 +116,12 @@ def _set_master_pass2(
if text == "/cancel": if text == "/cancel":
return _send_tmp_message(bot, mes.chat.id, "Успешная отмена") return _send_tmp_message(bot, mes.chat.id, "Успешная отмена")
password_hash, master_salt = encryption.master_pass.encrypt_master_pass( master_password = encryption.master_pass.encrypt_master_pass(
mes.from_user.id,
text, text,
) )
database.add.add_master_pass( database.add.add_master_pass(engine, master_password)
engine,
mes.from_user.id,
master_salt,
password_hash,
)
_send_tmp_message(bot, mes.chat.id, "Успех") _send_tmp_message(bot, mes.chat.id, "Успех")
del mes, text del mes, text
gc.collect() gc.collect()
@ -135,13 +133,16 @@ def reset_master_pass(
mes: Message, mes: Message,
) -> None: ) -> None:
_base_handler(bot, mes) _base_handler(bot, mes)
if database.get.get_master_pass(engine, mes.from_user.id) is None: if database.get.get_master_pass(engine, mes.from_user.id) is None:
return _send_tmp_message(bot, mes.chat.id, "Мастер пароль не задан") return _send_tmp_message(bot, mes.chat.id, "Мастер пароль не задан")
bot_mes = bot.send_message( bot_mes = bot.send_message(
mes.chat.id, mes.chat.id,
"Отправьте новый мастер пароль, осторожно, все текущие аккаунты " "Отправьте новый мастер пароль, осторожно, все текущие аккаунты "
"будут удалены навсегда", "будут удалены навсегда",
) )
bot.register_next_step_handler( bot.register_next_step_handler(
mes, functools.partial(_reset_master_pass2, bot, engine, bot_mes) mes, functools.partial(_reset_master_pass2, bot, engine, bot_mes)
) )
@ -155,9 +156,13 @@ def _reset_master_pass2(
if text == "/cancel": if text == "/cancel":
return _send_tmp_message(bot, mes.chat.id, "Успешная отмена") return _send_tmp_message(bot, mes.chat.id, "Успешная отмена")
hash_, salt = encryption.master_pass.encrypt_master_pass(text) master_password = encryption.master_pass.encrypt_master_pass(
mes.from_user.id,
text,
)
database.delete.purge_accounts(engine, mes.from_user.id) database.delete.purge_accounts(engine, mes.from_user.id)
database.change.change_master_pass(engine, mes.from_user.id, salt, hash_) database.change.change_master_pass(engine, master_password)
_send_tmp_message( _send_tmp_message(
bot, mes.chat.id, "Все ваши аккаунты удалены, а мастер пароль изменён" bot, mes.chat.id, "Все ваши аккаунты удалены, а мастер пароль изменён"
) )
@ -267,24 +272,30 @@ def _add_account5(
if text == "/cancel": if text == "/cancel":
return _send_tmp_message(bot, mes.chat.id, "Успешная отмена") return _send_tmp_message(bot, mes.chat.id, "Успешная отмена")
salt, hash_ = database.get.get_master_pass(engine, mes.from_user.id) master_password = database.get.get_master_pass(engine, mes.from_user.id)
if not encryption.master_pass.check_master_pass(text, hash_, salt): if not encryption.master_pass.check_master_pass(text, master_password):
return _send_tmp_message( return _send_tmp_message(
bot, bot,
mes.chat.id, mes.chat.id,
"Не подходит главный пароль", "Не подходит главный пароль",
) )
name, login, passwd = data["name"], data["login"], data["passwd"] # 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"],
)
enc_login, enc_pass, salt = encryption.other_accounts.encrypt( encrypted_account = encryption.other_accounts.encrypt(
login, account,
passwd,
text, text,
) )
result = database.add.add_account( result = database.add.add_account(
engine, mes.from_user.id, name, salt, enc_login, enc_pass engine,
encrypted_account,
) )
_send_tmp_message( _send_tmp_message(
@ -293,19 +304,19 @@ def _add_account5(
"Успех" if result else "Произошла не предвиденная ошибка", "Успех" if result else "Произошла не предвиденная ошибка",
) )
del data, name, login, passwd, enc_login del data, account
gc.collect() gc.collect()
def get_account(bot: telebot.TeleBot, engine: Engine, mes: Message) -> None: def get_account(bot: telebot.TeleBot, engine: Engine, mes: Message) -> None:
_base_handler(bot, mes) _base_handler(bot, mes)
bot_mes = bot.send_message(mes.chat.id, "Отправьте название аккаунта")
master_pass = database.get.get_master_pass(engine, mes.from_user.id) master_pass = database.get.get_master_pass(engine, mes.from_user.id)
if master_pass is None: if master_pass is None:
return _send_tmp_message(bot, mes.chat.id, "Нет мастер пароля") return _send_tmp_message(bot, mes.chat.id, "Нет мастер пароля")
bot_mes = bot.send_message(mes.chat.id, "Отправьте название аккаунта")
bot.register_next_step_handler( bot.register_next_step_handler(
mes, functools.partial(_get_account2, bot, engine, bot_mes) mes, functools.partial(_get_account2, bot, engine, bot_mes)
) )
@ -340,36 +351,28 @@ def _get_account3(
if text == "/cancel": if text == "/cancel":
return _send_tmp_message(bot, mes.chat.id, "Успешная отмена") return _send_tmp_message(bot, mes.chat.id, "Успешная отмена")
master_salt, hash_pass = database.get.get_master_pass( master_password = database.get.get_master_pass(
engine, engine,
mes.from_user.id, mes.from_user.id,
) )
if not encryption.master_pass.check_master_pass( if not encryption.master_pass.check_master_pass(text, master_password):
text,
hash_pass,
master_salt,
):
return _send_tmp_message(bot, mes.chat.id, "Не подходит мастер пароль") return _send_tmp_message(bot, mes.chat.id, "Не подходит мастер пароль")
salt, enc_login, enc_pass = database.get.get_account_info( account = database.get.get_account_info(engine, mes.from_user.id, name)
engine, mes.from_user.id, name account = encryption.other_accounts.decrypt(
) account,
login, passwd = encryption.other_accounts.decrypt(
enc_login,
enc_pass,
text, text,
salt,
) )
_send_tmp_message( _send_tmp_message(
bot, bot,
mes.chat.id, mes.chat.id,
f"Логин:\n`{login}`\nПароль:\n`{passwd}`\nНажмите на логин " f"Логин:\n`{account.login}`\nПароль:\n`{account.password}`\nНажмите "
"или пароль, чтобы скопировать", "на логин или пароль, чтобы скопировать",
30, 30,
) )
del text, mes, passwd, login del text, mes
gc.collect() gc.collect()
@ -455,15 +458,11 @@ def _export2(
if text == "/cancel": if text == "/cancel":
return _send_tmp_message(bot, mes.chat.id, "Успешная отмена") return _send_tmp_message(bot, mes.chat.id, "Успешная отмена")
master_salt, hash_pass = database.get.get_master_pass( master_password = database.get.get_master_pass(
engine, engine,
mes.from_user.id, mes.from_user.id,
) )
if not encryption.master_pass.check_master_pass( if not encryption.master_pass.check_master_pass(text, master_password):
text,
hash_pass,
master_salt,
):
return _send_tmp_message(bot, mes.chat.id, "Не подходит мастер пароль") return _send_tmp_message(bot, mes.chat.id, "Не подходит мастер пароль")
accounts = database.get.get_all_accounts(engine, mes.from_user.id) accounts = database.get.get_all_accounts(engine, mes.from_user.id)
@ -474,7 +473,7 @@ def _export2(
del text, accounts, json_io del text, accounts, json_io
gc.collect() gc.collect()
time.sleep(30) time.sleep(30)
bot.delete_message(bot_mes.chat.id, bot_mes.id) bot.delete_message(bot_mes.chat.id, bot_mes.id, timeout=5)
def import_accounts( def import_accounts(
@ -513,13 +512,16 @@ def _import2(
mes.chat.id, mes.chat.id,
"Вы должны отправить документ", "Вы должны отправить документ",
) )
if mes.document.file_size > 102_400: # If file size is bigger that 100 MB if mes.document.file_size > 102_400: # If file size is bigger than 100 MB
return _send_tmp_message(bot, mes.chat.id, "Файл слишком большой") return _send_tmp_message(bot, mes.chat.id, "Файл слишком большой")
file_info = bot.get_file(mes.document.file_id) file_info = bot.get_file(mes.document.file_id)
downloaded_file = bot.download_file(file_info.file_path) downloaded_file = bot.download_file(file_info.file_path)
try: try:
accounts = json_to_accounts(downloaded_file.decode("utf-8")) accounts = json_to_accounts(
downloaded_file.decode("utf-8"),
mes.from_user.id,
)
except Exception: except Exception:
return _send_tmp_message( return _send_tmp_message(
bot, bot,
@ -537,7 +539,7 @@ def _import3(
bot: telebot.TeleBot, bot: telebot.TeleBot,
engine: Engine, engine: Engine,
prev_mes: Message, prev_mes: Message,
accounts: list[tuple[str, str, str]], accounts: list[DecryptedAccount],
mes: Message, mes: Message,
) -> None: ) -> None:
_base_handler(bot, mes, prev_mes) _base_handler(bot, mes, prev_mes)
@ -545,40 +547,33 @@ def _import3(
if text == "/cancel": if text == "/cancel":
return _send_tmp_message(bot, mes.chat.id, "Успешная отмена") return _send_tmp_message(bot, mes.chat.id, "Успешная отмена")
master_salt, hash_pass = database.get.get_master_pass( master_password = database.get.get_master_pass(
engine, engine,
mes.from_user.id, mes.from_user.id,
) )
if not encryption.master_pass.check_master_pass( if not encryption.master_pass.check_master_pass(text, master_password):
text,
hash_pass,
master_salt,
):
return _send_tmp_message(bot, mes.chat.id, "Не подходит мастер пароль") return _send_tmp_message(bot, mes.chat.id, "Не подходит мастер пароль")
# 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] = []
for account in accounts: for account in accounts:
name, login, passwd = account if not check_account(account):
if not check_account(name, login, passwd): failed.append(account.name)
failed.append(name)
continue continue
enc_login, enc_passwd, salt = encryption.other_accounts.encrypt( account = encryption.other_accounts.encrypt(
login, account,
passwd,
text, text,
) )
result = database.add.add_account( result = database.add.add_account(engine, account)
engine, mes.from_user.id, name, salt, enc_login, enc_passwd
)
if not result: if not result:
failed.append(name) failed.append(account.name)
if failed: if failed:
mes_text = "Не удалось добавить:\n" + "\n".join(failed) mes_text = "Не удалось добавить:\n" + "\n".join(failed)
else: else:
mes_text = "Успех" mes_text = "Успех"
_send_tmp_message(bot, mes.chat.id, mes_text, 10) _send_tmp_message(bot, mes.chat.id, mes_text, 10)
del text, mes, accounts del text, mes, accounts
gc.collect() gc.collect()

21
src/classes.py Normal file
View File

@ -0,0 +1,21 @@
from typing import Self, TypeAlias
import pydantic
from .database import models
Account: TypeAlias = models.Account
class DecryptedAccount(pydantic.BaseModel):
user_id: int
name: str
login: str
password: str
@classmethod
def from_tuple(cls, tuple_: tuple[str, str, str]) -> Self:
return cls(name=tuple_[0], login=tuple_[1], password=tuple_[2])
def as_tuple(self) -> tuple[str, str, str]:
return (self.name, self.login, self.password)

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"] __all__ = ["add", "delete", "get", "models", "prepare", "change"]

View File

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

View File

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

View File

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

View File

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

View File

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