Danbooru Prompt Selector/TEST2024/NAIA_Login.py

from base64 import urlsafe_b64encode
from hashlib import blake2b
import argon2
import requests

BASE_URL="https://api.novelai.net"

def argon_hash(email: str, password: str, size: int, domain: str) -> str:
    pre_salt = f"{password[:6]}{email}{domain}"
    # salt
    blake = blake2b(digest_size=16)
    blake.update(pre_salt.encode())
    salt = blake.digest()
    raw = argon2.low_level.hash_secret_raw(
        password.encode(),
        salt,
        2,
        int(2000000 / 1024),
        1,
        size,
        argon2.low_level.Type.ID,
    )
    hashed = urlsafe_b64encode(raw).decode()
    return hashed

def login(key) -> str:
    response = requests.post(f"{BASE_URL}/user/login", json={ "key": key })
    # catch any errors
    return response.json()["accessToken"]

def get_access_key(email: str, password: str) -> str:
    return argon_hash(email, password, 64, "novelai_data_access_key")[:64]