File: signed_token.py

package info (click to toggle)
mautrix-python 0.20.7-1
  • links: PTS, VCS
  • area: main
  • in suites: trixie
  • size: 1,812 kB
  • sloc: python: 19,103; makefile: 16
file content (45 lines) | stat: -rw-r--r-- 1,334 bytes parent folder | download
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
# Copyright (c) 2022 Tulir Asokan
#
# This Source Code Form is subject to the terms of the Mozilla Public
# License, v. 2.0. If a copy of the MPL was not distributed with this
# file, You can obtain one at http://mozilla.org/MPL/2.0/.
from __future__ import annotations

from hashlib import sha256
import base64
import hmac
import json


def _get_checksum(key: str, payload: bytes) -> str:
    hasher = hmac.new(key.encode("utf-8"), msg=payload, digestmod=sha256)
    checksum = base64.urlsafe_b64encode(hasher.digest())
    return checksum.decode("utf-8").rstrip("=")


def sign_token(key: str, payload: dict) -> str:
    payload_b64 = base64.urlsafe_b64encode(json.dumps(payload).encode("utf-8"))
    checksum = _get_checksum(key, payload_b64)
    payload_str = payload_b64.decode("utf-8").rstrip("=")
    return f"{checksum}:{payload_str}"


def verify_token(key: str, data: str) -> dict | None:
    if not data:
        return None

    try:
        checksum, payload = data.split(":", 1)
    except ValueError:
        return None

    payload += (3 - (len(payload) + 3) % 4) * "="

    if checksum != _get_checksum(key, payload.encode("utf-8")):
        return None

    payload = base64.urlsafe_b64decode(payload).decode("utf-8")
    try:
        return json.loads(payload)
    except json.JSONDecodeError:
        return None