import base64
import hmac
import json
import secrets


def gen_token(key: bytes, data) -> str:
    payload = json.dumps(data).encode()
    mac = base64.urlsafe_b64encode(hmac.new(key, payload, "sha256").digest())
    return base64.urlsafe_b64encode(payload).decode() + "." + mac.decode()


def get_data(key: bytes, token: str) -> Any:
    payload, mac = token.split(".", 1)
    payload_bytes = base64.urlsafe_b64decode(payload)
    mac_bytes = base64.urlsafe_b64decode(mac)
    expected_mac = hmac.new(key, payload_bytes, "sha256").digest()
    if not hmac.compare_digest(mac_bytes, expected_mac):
        raise ValueError("Invalid signature")
    return json.loads(payload_bytes.decode())


if __name__ == "__main__":
    key = secrets.token_bytes(32)
    data = {"some": {"nested": "data"}}
    token = gen_token(key, data)
    print(token)
    print(get_data(key, token))

# Prints out something like this:
# eyJzb21lIjogeyJuZXN0ZWQiOiAiZGF0YSJ9fQ==.KKnjFS9xHPqi6uVHDAH-L52FkjH05DHKi0QZxkMwN7w=
# {'some': {'nested': 'data'}}