tokens.py

"""

Generating a random token::

	>>> random_token() 		# doctest: +SKIP
	10239537334670678660

Signing Tokens (internally, this uses a truncated HMAC)::

	>>> key = b'secret'
	>>> signed_token = sign(1234567890, key, 64)
	>>> signed_token
	22773757911674290731578468199

	>>> unsign(signed_token, key, 64)
	1234567890

Expressing Tokens in base-62::

	>>> encode(123456789)
	'8m0Kx'

	>>> decode('8m0Kx')
	123456789

"""

import secrets
import hmac

from string import ascii_uppercase, ascii_letters, digits

_DIGITS = digits + ascii_letters

def encode(number, alphabet=_DIGITS, base=None, min_length=1):
	"""
	Example::

		>>> encode(0)
		'0'

		>>> encode(123456789012)
		'2aL26ws'

		>>> encode(1234, min_length=8)
		'000000jU'

	"""
	base = base or len(alphabet)

	if number < 0:
		raise ValueError
	result = []

	while number:
		number, digit = divmod(number, base)
		result.append(_DIGITS[digit])

	result.extend([_DIGITS[0]] * (min_length - len(result)))

	return ''.join(reversed(result))

def decode(chars, alphabet=_DIGITS, base=None):
	"""
	Examples::

		>>> decode('2aL26ws')
		123456789012

		>>> decode('0000xyz') == decode('xyz')
		True



	"""
	base = base or len(alphabet)
	number = 0
	l = len(_DIGITS)
	for c in chars:
		number *= base
		number += _DIGITS.index(c)
	return number

class SignatureFailedToVerify(Exception):
	pass

def random_token(bits=64):
	return secrets.randbits(bits)

def sign(tok, key, signature_bits):
	"""
	Examples::
		>>> sign(123, b'secret', 64)
		2275831507231871192277
	"""
	return (tok << signature_bits) + signature(tok, key, signature_bits)

def signature(tok, key, signature_bits):
	"""
	Examples::

		>>> signature(1, b'secret', 64)
		2422091285121405787
	"""
	mac = hmac.new(key, bytes(str(int(tok)), "ascii"), "sha256").digest()
	mac = int.from_bytes(mac, "big")
	return mac % (1 << signature_bits)

def unsign(tok, key, signature_bits):
	"""
		>>> spotless = 2275831507231871192277
		>>> unsign(spotless, b'secret', 64)
		123

		>>> less_than_spotless = spotless - 1
		>>> unsign(less_than_spotless, b'secret', 64)
	 	Traceback (most recent call last):
		...
		tokens.SignatureFailedToVerify

		>>> unsign(less_than_spotless, None, 64)
		123

	"""
	result = tok >> signature_bits
	if key is not None:
		sig_provided = str(tok % (1 << signature_bits))
		sig_calculated = str(signature(result, key, signature_bits))
		if not hmac.compare_digest(sig_provided, sig_calculated):	
			raise SignatureFailedToVerify
	return result