src/__init__.py

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

src/cryptography/other_accounts.py

@@ -8,7 +8,7 @@ from cryptography.hazmat.primitives import hashes
 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(
         algorithm=hashes.SHA256(),
         length=32,
@@ -16,17 +16,17 @@ def _generate_key(salt: bytes, master_pass: bytes) -> bytes:
         iterations=100000,
         backend=default_backend(),
     )
-    key = base64.urlsafe_b64encode(kdf.derive(master_pass))
+    key = base64.urlsafe_b64encode(kdf.derive(master_pass_hash))
     return key
 
 
 def encrypt_account_info(
-    login: str, passwd: str, master_pass: bytes
+    login: str, passwd: str, master_pass_hash: bytes
 ) -> tuple[bytes, bytes, bytes]:
     """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"""
     salt = bcrypt.gensalt()
-    key = _generate_key(salt, master_pass)
+    key = _generate_key(salt, master_pass_hash)
     f = Fernet(key)
     enc_login = f.encrypt(login.encode("utf-8"))
     enc_passwd = f.encrypt(passwd.encode("utf-8"))
@@ -34,9 +34,9 @@ def encrypt_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]:
-    key = _generate_key(salt, master_pass)
+    key = _generate_key(salt, master_pass_hash)
     f = Fernet(key)
     login_bytes = f.decrypt(enc_login)
     pass_bytes = f.decrypt(enc_pass)

src/database/__init__.py

@@ -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"]

src/database/add.py

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

src/database/delete.py

@@ -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()

src/database/get.py

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

src/database/models.py

@@ -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)
-    )

src/database/prepare.py

@@ -1,16 +1,35 @@
-import sqlmodel
-from sqlalchemy.future import Engine
-
-from . import models
+import mariadb
 
 
-def get_engine(host: str, user: str, passwd: str, db: str) -> Engine:
-    engine = sqlmodel.create_engine(
-        f"mariadb+mariadbconnector://{user}:{passwd}@{host}/{db}"
+def _create_tables(con: mariadb.Connection) -> None:
+    cursor = con.cursor()
+    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))
-    return engine
+    cursor.execute(
+        """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:
-    sqlmodel.SQLModel.metadata.create_all(engine, [models.Account, models.MasterPass])
+def _create_index(con: mariadb.Connection) -> None:
+    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)