utils.py

from datetime import datetime
from zoneinfo import ZoneInfo
from typing import Any
from jose import jwt
import hashlib
import json
import uuid
import os


SECRET_KEY = os.environ.get("SECRET_KEY")


def beijing():
    """get beijing time"""
    return datetime.now(tz=ZoneInfo("Asia/Shanghai"))


def decode_jwt(token: str) -> dict[str, Any]:
    """get payload in the jwt token"""
    assert SECRET_KEY, "Please set the environment variable SECRET_KEY"
    return jwt.decode(token, SECRET_KEY, algorithms=["HS256"])


def parse_token(token: str | None) -> tuple[str, str]:
    """parse user token to get uid and username

    :param token (str): jwt token
    :return (str, str): (uid, username)
    """
    if not token:
        uid, username = "no_uid", "Anonymous"
    else:
        payload: dict = decode_jwt(token)
        uid = payload.get("uid", "no_uid")
        username = payload.get("username", "Anonymous")
    print(f"UID: {uid}")
    print(f"Username: {username}")
    return uid, username


def md5(text: list[str | bytes] | str | bytes | None = None) -> str:
    """generate the md5 hash code of the given text, if text is None,
    return a random md5"""
    code = hashlib.md5()
    if text:
        if isinstance(text, str):
            text = text.encode("utf-8")
            code.update(text)
        elif isinstance(text, list):
            for t in text:
                if isinstance(t, str):
                    t = t.encode("utf-8")
                code.update(t)
        else:
            code.update(text)
    else:
        code.update(uuid.uuid4().bytes)
    return code.hexdigest()


def json_to_str(obj: dict | list) -> str:
    return json.dumps(obj, separators=(",", ":"))