X-Financial/server/src/app/core/security.py

from __future__ import annotations

import hashlib
import secrets
from base64 import urlsafe_b64decode, urlsafe_b64encode

PBKDF2_ALGORITHM = "sha256"
PBKDF2_ITERATIONS = 120_000
SALT_BYTES = 16


def hash_password(password: str) -> str:
    salt = secrets.token_bytes(SALT_BYTES)
    digest = hashlib.pbkdf2_hmac(
        PBKDF2_ALGORITHM,
        password.encode("utf-8"),
        salt,
        PBKDF2_ITERATIONS,
    )
    encoded_salt = urlsafe_b64encode(salt).decode("utf-8")
    encoded_digest = urlsafe_b64encode(digest).decode("utf-8")
    return f"pbkdf2_{PBKDF2_ALGORITHM}${PBKDF2_ITERATIONS}${encoded_salt}${encoded_digest}"


def verify_password(password: str, password_hash: str) -> bool:
    if password_hash.startswith("scrypt$"):
        return verify_scrypt_password(password, password_hash)

    try:
        scheme, iterations, encoded_salt, encoded_digest = password_hash.split("$", 3)
    except ValueError:
        return False

    if scheme != f"pbkdf2_{PBKDF2_ALGORITHM}":
        return False

    salt = urlsafe_b64decode(encoded_salt.encode("utf-8"))
    expected_digest = urlsafe_b64decode(encoded_digest.encode("utf-8"))
    computed_digest = hashlib.pbkdf2_hmac(
        PBKDF2_ALGORITHM,
        password.encode("utf-8"),
        salt,
        int(iterations),
    )
    return secrets.compare_digest(computed_digest, expected_digest)


def verify_scrypt_password(password: str, password_hash: str) -> bool:
    try:
        scheme, n_value, r_value, p_value, key_length, salt_hex, derived_key_hex = password_hash.split("$", 6)
    except ValueError:
        return False

    if scheme != "scrypt":
        return False

    try:
        salt = bytes.fromhex(salt_hex)
        expected_key = bytes.fromhex(derived_key_hex)
        derived_key = hashlib.scrypt(
            password.encode("utf-8"),
            salt=salt,
            n=int(n_value),
            r=int(r_value),
            p=int(p_value),
            dklen=int(key_length),
        )
    except ValueError:
        return False

    return secrets.compare_digest(derived_key, expected_key)