database.py

import os
import sys
import json
import csv
import sqlite3
import base64
from datetime import datetime
from cryptography.fernet import Fernet, InvalidToken
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC

# Locate the database file next to the main executable or script
if getattr(sys, 'frozen', False):
    executable_dir = os.path.dirname(sys.executable)
else:
    executable_dir = os.path.dirname(os.path.abspath(__file__))

DB_FILE = os.path.join(executable_dir, "vault.db")
VERIFICATION_STRING = b"verification_token"

def get_db_connection():
    """Establishes and returns a connection to the SQLite database."""
    conn = sqlite3.connect(DB_FILE)
    conn.row_factory = sqlite3.Row
    return conn

def db_init():
    """Initializes the SQLite database, creates tables, and handles dynamic migrations."""
    conn = get_db_connection()
    cursor = conn.cursor()
    
    # Settings table: for vault initialization details (salt, verifier token)
    cursor.execute('''
        CREATE TABLE IF NOT EXISTS settings (
            key TEXT PRIMARY KEY,
            val BLOB NOT NULL
        )
    ''')
    
    # Credentials table: for storing encrypted password details
    cursor.execute('''
        CREATE TABLE IF NOT EXISTS credentials (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            service TEXT NOT NULL,
            username TEXT NOT NULL,
            password_encrypted BLOB NOT NULL
        )
    ''')
    
    # Dynamic Migration: Check existing columns in credentials table
    cursor.execute("PRAGMA table_info(credentials)")
    existing_columns = {column['name'] for column in cursor.fetchall()}
    
    # Add new columns if missing
    new_columns = [
        ("url", "TEXT DEFAULT ''"),
        ("notes", "TEXT DEFAULT ''"),
        ("category", "TEXT DEFAULT 'General'"),
        ("totp_secret_encrypted", "BLOB DEFAULT NULL"),
        ("created_at", "TIMESTAMP DEFAULT CURRENT_TIMESTAMP"),
        ("updated_at", "TIMESTAMP DEFAULT CURRENT_TIMESTAMP")
    ]
    
    for col_name, col_type in new_columns:
        if col_name not in existing_columns:
            cursor.execute(f"ALTER TABLE credentials ADD COLUMN {col_name} {col_type}")
            
    conn.commit()
    conn.close()

def is_vault_initialized() -> bool:
    """Checks whether the vault has been set up with a master password."""
    db_init()
    conn = get_db_connection()
    cursor = conn.cursor()
    cursor.execute("SELECT val FROM settings WHERE key = 'master_verifier'")
    row = cursor.fetchone()
    conn.close()
    return row is not None

def derive_session_key(master_password: str, salt: bytes) -> bytes:
    """
    Derives a 32-byte Fernet key from the master password using PBKDF2.
    Using a high number of iterations (100,000) prevents brute-forcing offline database dumps.
    """
    kdf = PBKDF2HMAC(
        algorithm=hashes.SHA256(),
        length=32,
        salt=salt,
        iterations=100000
    )
    return base64.urlsafe_b64encode(kdf.derive(master_password.encode('utf-8')))

def initialize_vault(master_password: str):
    """
    Sets up the vault for the first time by generating a master salt and
    storing an encrypted verification token.
    """
    db_init()
    conn = get_db_connection()
    cursor = conn.cursor()
    
    master_salt = os.urandom(16)
    key = derive_session_key(master_password, master_salt)
    
    fernet = Fernet(key)
    verifier_encrypted = fernet.encrypt(VERIFICATION_STRING)
    
    cursor.execute("INSERT OR REPLACE INTO settings (key, val) VALUES (?, ?)", ("master_salt", master_salt))
    cursor.execute("INSERT OR REPLACE INTO settings (key, val) VALUES (?, ?)", ("master_verifier", verifier_encrypted))
    
    conn.commit()
    conn.close()

def verify_master_password(master_password: str) -> bytes | None:
    """
    Verifies if the master password is correct.
    Returns the derived session key (bytes) on success, or None on failure.
    """
    if not is_vault_initialized():
        return None
        
    conn = get_db_connection()
    cursor = conn.cursor()
    
    cursor.execute("SELECT val FROM settings WHERE key = 'master_salt'")
    salt_row = cursor.fetchone()
    cursor.execute("SELECT val FROM settings WHERE key = 'master_verifier'")
    verifier_row = cursor.fetchone()
    conn.close()
    
    if not salt_row or not verifier_row:
        return None
        
    salt = salt_row['val']
    verifier_encrypted = verifier_row['val']
    
    key = derive_session_key(master_password, salt)
    
    try:
        fernet = Fernet(key)
        decrypted = fernet.decrypt(verifier_encrypted)
        if decrypted == VERIFICATION_STRING:
            return key
    except (InvalidToken, ValueError):
        pass
        
    return None

def add_credential(session_key: bytes, service: str, username: str, password_plaintext: str,
                   url: str = "", notes: str = "", category: str = "General", totp_secret: str = ""):
    """Encrypts the credential password and TOTP secret, then saves to the database."""
    fernet = Fernet(session_key)
    encrypted_password = fernet.encrypt(password_plaintext.encode('utf-8'))
    
    encrypted_totp = None
    if totp_secret:
        encrypted_totp = fernet.encrypt(totp_secret.strip().encode('utf-8'))
        
    now_str = datetime.now().isoformat()
    
    conn = get_db_connection()
    cursor = conn.cursor()
    cursor.execute(
        '''INSERT INTO credentials 
           (service, username, password_encrypted, url, notes, category, totp_secret_encrypted, created_at, updated_at) 
           VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)''',
        (service, username, encrypted_password, url, notes, category, encrypted_totp, now_str, now_str)
    )
    conn.commit()
    conn.close()

def update_credential(session_key: bytes, cred_id: int, service: str, username: str, password_plaintext: str,
                      url: str = "", notes: str = "", category: str = "General", totp_secret: str = ""):
    """Updates an existing credential in the database."""
    fernet = Fernet(session_key)
    encrypted_password = fernet.encrypt(password_plaintext.encode('utf-8'))
    
    encrypted_totp = None
    if totp_secret:
        encrypted_totp = fernet.encrypt(totp_secret.strip().encode('utf-8'))
        
    now_str = datetime.now().isoformat()
    
    conn = get_db_connection()
    cursor = conn.cursor()
    cursor.execute(
        '''UPDATE credentials 
           SET service = ?, username = ?, password_encrypted = ?, url = ?, notes = ?, category = ?, 
               totp_secret_encrypted = ?, updated_at = ?
           WHERE id = ?''',
        (service, username, encrypted_password, url, notes, category, encrypted_totp, now_str, cred_id)
    )
    conn.commit()
    conn.close()

def get_all_credentials(session_key: bytes, search_query: str = "", category_filter: str = "All") -> list:
    """
    Fetches all credentials from the DB, decrypts passwords/TOTP keys,
    and returns a list of dictionaries. Filters by query or category if provided.
    """
    conn = get_db_connection()
    cursor = conn.cursor()
    
    query = "SELECT id, service, username, password_encrypted, url, notes, category, totp_secret_encrypted, created_at, updated_at FROM credentials WHERE 1=1"
    params = []
    
    if search_query:
        query += " AND (service LIKE ? OR username LIKE ? OR notes LIKE ? OR url LIKE ?)"
        term = f"%{search_query}%"
        params.extend([term, term, term, term])
        
    if category_filter and category_filter != "All":
        query += " AND category = ?"
        params.append(category_filter)
        
    query += " ORDER BY service ASC"
    
    cursor.execute(query, params)
    rows = cursor.fetchall()
    conn.close()
    
    fernet = Fernet(session_key)
    results = []
    
    for row in rows:
        try:
            decrypted_password = fernet.decrypt(row['password_encrypted']).decode('utf-8')
            
            totp_secret = ""
            if row['totp_secret_encrypted']:
                try:
                    totp_secret = fernet.decrypt(row['totp_secret_encrypted']).decode('utf-8')
                except (InvalidToken, ValueError):
                    totp_secret = ""
                    
            results.append({
                "id": row['id'],
                "service": row['service'],
                "username": row['username'],
                "password": decrypted_password,
                "url": row['url'] or "",
                "notes": row['notes'] or "",
                "category": row['category'] or "General",
                "totp_secret": totp_secret,
                "created_at": row['created_at'] or "",
                "updated_at": row['updated_at'] or ""
            })
        except (InvalidToken, ValueError):
            continue
            
    return results

def delete_credential(cred_id: int):
    """Deletes a credential from the vault by database ID."""
    conn = get_db_connection()
    cursor = conn.cursor()
    cursor.execute("DELETE FROM credentials WHERE id = ?", (cred_id,))
    conn.commit()
    conn.close()

def export_vault_encrypted(session_key: bytes, file_path: str):
    """Exports all credentials to a Fernet-encrypted JSON file."""
    creds = get_all_credentials(session_key)
    raw_data = json.dumps(creds).encode('utf-8')
    
    fernet = Fernet(session_key)
    encrypted_data = fernet.encrypt(raw_data)
    
    with open(file_path, 'wb') as f:
        f.write(encrypted_data)

def import_vault_encrypted(session_key: bytes, file_path: str) -> int:
    """Decrypts and imports credentials from an encrypted .spvault backup file. Returns count of imported items."""
    with open(file_path, 'rb') as f:
        encrypted_data = f.read()
        
    fernet = Fernet(session_key)
    decrypted_data = fernet.decrypt(encrypted_data)
    creds = json.loads(decrypted_data.decode('utf-8'))
    
    count = 0
    for item in creds:
        add_credential(
            session_key,
            service=item.get("service", "Imported Service"),
            username=item.get("username", ""),
            password_plaintext=item.get("password", ""),
            url=item.get("url", ""),
            notes=item.get("notes", ""),
            category=item.get("category", "General"),
            totp_secret=item.get("totp_secret", "")
        )
        count += 1
    return count

def export_vault_csv(session_key: bytes, file_path: str):
    """Exports all credentials to a plain text CSV file."""
    creds = get_all_credentials(session_key)
    fieldnames = ["service", "username", "password", "url", "category", "totp_secret", "notes"]
    
    with open(file_path, 'w', newline='', encoding='utf-8') as f:
        writer = csv.DictWriter(f, fieldnames=fieldnames)
        writer.writeheader()
        for c in creds:
            writer.writerow({
                "service": c["service"],
                "username": c["username"],
                "password": c["password"],
                "url": c["url"],
                "category": c["category"],
                "totp_secret": c["totp_secret"],
                "notes": c["notes"]
            })

def import_vault_csv(session_key: bytes, file_path: str) -> int:
    """Imports credentials from a CSV file (supports standard column headers). Returns imported count."""
    count = 0
    with open(file_path, 'r', encoding='utf-8') as f:
        reader = csv.DictReader(f)
        for row in reader:
            service = row.get("service") or row.get("name") or row.get("title") or "Imported"
            username = row.get("username") or row.get("login_username") or row.get("email") or ""
            password = row.get("password") or row.get("login_password") or ""
            url = row.get("url") or row.get("login_uri") or ""
            notes = row.get("notes") or row.get("comment") or ""
            category = row.get("category") or "General"
            totp_secret = row.get("totp_secret") or row.get("totp") or ""
            
            if service and password:
                add_credential(
                    session_key,
                    service=service,
                    username=username,
                    password_plaintext=password,
                    url=url,
                    notes=notes,
                    category=category,
                    totp_secret=totp_secret
                )
                count += 1
    return count