Compare commits

..

No commits in common. "29b7125a4dcf0593b85cdd67943613a64f0f8f86" and "e09c56759eb81fa5e7d9b4821f734b99c7a915a9" have entirely different histories.

8 changed files with 114 additions and 122 deletions

View File

@ -1,24 +1,20 @@
import os import os
import mariadb
from dotenv import load_dotenv from dotenv import load_dotenv
from sqlalchemy.future import Engine
from . import bot, cryptography, database from . import bot, cryptography, database
__all__ = ["bot", "cryptography", "database"] __all__ = ["bot", "cryptography", "database"]
engine: Engine
def main() -> None: def main() -> None:
global engine
load_dotenv("./.env") load_dotenv("./.env")
engine = database.prepare.get_engine( con = mariadb.connect(
host=os.getenv("DB_HOST"), os.getenv("DB_HOST"),
user=os.getenv("DB_USER"), os.getenv("DB_USER"),
passwd=os.getenv("DB_PASS"), os.getenv("DB_PASS"),
db=os.getenv("DB_NAME"), os.getenv("DB_NAME"),
) # type: ignore )
database.prepare.prepare(engine) database.prepare(con)
bot_ = bot.create_bot(os.getenv("TG_TOKEN"), con) # type: ignore
bot_.infinity_polling()

View File

@ -8,7 +8,7 @@ from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
def _generate_key(salt: bytes, master_pass: bytes) -> bytes: def _generate_key(salt: bytes, master_pass_hash: bytes) -> bytes:
kdf = PBKDF2HMAC( kdf = PBKDF2HMAC(
algorithm=hashes.SHA256(), algorithm=hashes.SHA256(),
length=32, length=32,
@ -16,17 +16,17 @@ def _generate_key(salt: bytes, master_pass: bytes) -> bytes:
iterations=100000, iterations=100000,
backend=default_backend(), backend=default_backend(),
) )
key = base64.urlsafe_b64encode(kdf.derive(master_pass)) key = base64.urlsafe_b64encode(kdf.derive(master_pass_hash))
return key return key
def encrypt_account_info( def encrypt_account_info(
login: str, passwd: str, master_pass: bytes login: str, passwd: str, master_pass_hash: bytes
) -> tuple[bytes, bytes, bytes]: ) -> tuple[bytes, bytes, bytes]:
"""Encrypts login and password of a user using hash of their master password as a key. """Encrypts login and password of a user using hash of their master password as a key.
Returns a tuple of encrypted login password and salt""" Returns a tuple of encrypted login password and salt"""
salt = bcrypt.gensalt() salt = bcrypt.gensalt()
key = _generate_key(salt, master_pass) key = _generate_key(salt, master_pass_hash)
f = Fernet(key) f = Fernet(key)
enc_login = f.encrypt(login.encode("utf-8")) enc_login = f.encrypt(login.encode("utf-8"))
enc_passwd = f.encrypt(passwd.encode("utf-8")) enc_passwd = f.encrypt(passwd.encode("utf-8"))
@ -34,9 +34,9 @@ def encrypt_account_info(
def decrypt_account_info( def decrypt_account_info(
enc_login: bytes, enc_pass: bytes, master_pass: bytes, salt: bytes enc_login: bytes, enc_pass: bytes, master_pass_hash: bytes, salt: bytes
) -> tuple[str, str]: ) -> tuple[str, str]:
key = _generate_key(salt, master_pass) key = _generate_key(salt, master_pass_hash)
f = Fernet(key) f = Fernet(key)
login_bytes = f.decrypt(enc_login) login_bytes = f.decrypt(enc_login)
pass_bytes = f.decrypt(enc_pass) pass_bytes = f.decrypt(enc_pass)

View File

@ -1,3 +1,3 @@
from . import add, delete, get, models, prepare from . import add, delete, get, prepare
__all__ = ["add", "delete", "get", "models", "prepare"] __all__ = ["add", "delete", "get", "prepare"]

View File

@ -1,40 +1,42 @@
import sqlmodel import traceback
import mariadb
from sqlalchemy.future import Engine
from . import models import mariadb
def add_master_pass(
id: int, hashed_passwd: bytes, salt: bytes, con: mariadb.Connection
) -> bool:
cursor = con.cursor()
try:
cursor.execute(
"INSERT INTO master_pass (user_id, salt, passwd) VALUES (?, ?, ?)",
[id, hashed_passwd, salt],
)
cursor.close()
except Exception:
traceback.print_exc()
return False
else:
return True
def add_account( def add_account(
engine: Engine, id: int,
user_id: int, acc_name: str,
name: str,
salt: bytes, salt: bytes,
enc_login: bytes, enc_login: bytes,
enc_pass: bytes, enc_passwd: bytes,
con: mariadb.Connection,
) -> bool: ) -> bool:
"""Adds account to db. Returns true, if on success""" cursor = con.cursor()
account = models.Account(
user_id=user_id, name=name, salt=salt, enc_login=enc_login, enc_pass=enc_pass
)
try: try:
with sqlmodel.Session(engine) as session: cursor.execute(
session.add(account) "INSERT INTO accounts (user_id, acc_name, salt, enc_login, enc_pass) VALUES (?, ?, ?, ?, ?, ?)",
session.commit() [id, acc_name, salt, enc_login, enc_passwd],
except Exception: )
return False cursor.close()
else:
return True
def add_master_pass(engine: Engine, user_id: int, salt: bytes, passwd: bytes) -> bool:
"""Adds master password to db. Returns true, if on success"""
master_pass = models.MasterPass(user_id=user_id, salt=salt, passwd=passwd)
try:
with sqlmodel.Session(engine) as session:
session.add(master_pass)
session.commit()
except Exception: except Exception:
traceback.print_exc()
return False return False
else: else:
return True return True

View File

@ -0,0 +1,15 @@
import mariadb
def delete_master_pass(con: mariadb.Connection, user_id: int) -> None:
cursor = con.cursor()
cursor.execute("DELETE FROM master_pass WHERE user_id=?", [user_id])
cursor.close()
def delete_account(con: mariadb.Connection, user_id: int, account: str):
cursor = con.cursor()
cursor.execute(
"DELETE FROM accounts WHERE user_id = ? AND acc_name = ?", [user_id, account]
)
cursor.close()

View File

@ -1,39 +1,31 @@
import sqlmodel import mariadb
from sqlalchemy.future import Engine
from . import models
def get_master_pass(engine: Engine, user_id: int) -> tuple[bytes, bytes] | None: def get_master_pass(con: mariadb.Connection, id: int) -> tuple[bytes, bytes]:
"""Gets master pass. Returns tuple of salt and password""" """Returns tuple of salt and hashed master password"""
statement = sqlmodel.select(models.MasterPass).where( cursor = con.cursor()
models.MasterPass.user_id == user_id cursor.execute("SELECT salt, passwd FROM master_pass IF user_id = ?", [id])
) result = cursor.fetchone()
with sqlmodel.Session(engine) as session: cursor.close()
result = session.exec(statement).first() return result
print(result)
if result is None:
return
return (result.salt, result.passwd)
def get_accounts(engine: Engine, user_id: int) -> list[str]: def get_accounts(con: mariadb.Connection, id: int) -> list[str]:
"""Gets list of account names""" """Returns list of user accounts"""
statement = sqlmodel.select(models.Account).where(models.Account.user_id == user_id) cursor = con.cursor()
with sqlmodel.Session(engine) as session: cursor.execute("SELECT acc_name FROM accounts IF user_id = ?", [id])
result = session.exec(statement) return [i[0] for i in cursor.fetchall()]
return [account.name for account in result]
def get_account_info( def get_account_info(
engine: Engine, user_id: int, name: str id: int, name: str, con: mariadb.Connection
) -> tuple[bytes, bytes, bytes]: ) -> tuple[bytes, bytes, bytes]:
"""Gets account info. Returns tuple of salt, login and password""" """Returns tuple of salt, login and password"""
statement = sqlmodel.select(models.Account).where( cursor = con.cursor()
models.Account.user_id == user_id and models.Account.name == name cursor.execute(
"""SELECT salt, enc_login, enc_pass IF user_id = ? AND acc_name = ?""",
[id, name],
) )
with sqlmodel.Session(engine) as session: result = cursor.fetchone()
result = session.exec(statement).first() cursor.close()
if result is None: return result
return
return (result.salt, result.enc_login, result.enc_pass)

View File

@ -1,32 +0,0 @@
from typing import Optional
import sqlmodel
class MasterPass(sqlmodel.SQLModel, table=True):
__tablename__ = "master_passwords"
id: Optional[int] = sqlmodel.Field(primary_key=True)
user_id: int = sqlmodel.Field(nullable=False, index=True, unique=True)
salt: bytes = sqlmodel.Field(
sa_column=sqlmodel.Column(type_=sqlmodel.VARBINARY(255), nullable=False)
)
passwd: bytes = sqlmodel.Field(
sa_column=sqlmodel.Column(type_=sqlmodel.VARBINARY(255), nullable=False)
)
class Account(sqlmodel.SQLModel, table=True):
__tablename__ = "accounts"
__table_args__ = (sqlmodel.UniqueConstraint("user_id", "name"),)
id: Optional[int] = sqlmodel.Field(primary_key=True)
user_id: int = sqlmodel.Field(nullable=False, index=True)
name: str = sqlmodel.Field(nullable=False, index=True, max_length=255)
salt: bytes = sqlmodel.Field(
sa_column=sqlmodel.Column(type_=sqlmodel.VARBINARY(255), nullable=False)
)
enc_login: bytes = sqlmodel.Field(
sa_column=sqlmodel.Column(type_=sqlmodel.VARBINARY(255), nullable=False)
)
enc_pass: bytes = sqlmodel.Field(
sa_column=sqlmodel.Column(type_=sqlmodel.VARBINARY(255), nullable=False)
)

View File

@ -1,16 +1,35 @@
import sqlmodel import mariadb
from sqlalchemy.future import Engine
from . import models
def get_engine(host: str, user: str, passwd: str, db: str) -> Engine: def _create_tables(con: mariadb.Connection) -> None:
engine = sqlmodel.create_engine( cursor = con.cursor()
f"mariadb+mariadbconnector://{user}:{passwd}@{host}/{db}" cursor.execute(
"""CREATE TABLE IF NOT EXISTS master_pass (user_id INT,
salt BINARY(64),
passwd BINARY(64),
PRIMARY KEY(user_id)
)"""
) )
print(type(engine)) cursor.execute(
return engine """CREATE TABLE IF NOT EXISTS accounts(user_id INT,
acc_name VARCHAR(255),
salt BINARY(64),
enc_login BINARY(64),
enc_pass BINARY(64),
UNIQUE(acc_name, user_id)
)"""
)
cursor.close()
def prepare(engine: Engine) -> None: def _create_index(con: mariadb.Connection) -> None:
sqlmodel.SQLModel.metadata.create_all(engine, [models.Account, models.MasterPass]) cursor = con.cursor()
cursor.execute(
"""CREATE INDEX IF NOT EXISTS user_id_to_acc on accounts(user_id)
"""
)
def prepare(con: mariadb.Connection) -> None:
_create_tables(con)
_create_index(con)