# deck.otp: HOTP/TOTP implementation for T-DeckARD
# Initially based on the utotp code
# Should be compatible with CircuitPython, MicroPython and CPython
# On CircuitPython, depends on adafruit_hashlib
# Usage: from deck.otp import hotp, totp
# Created by Luxferre in 2025, released into public domain

try:
  import ustruct as struct
except ImportError:
  import struct

try:
  from adafruit_hashlib import sha1
except ImportError:
  from hashlib import sha1

FORCE_MPY = False
try:
  from deck.time import datetime, timedelta
except ImportError:
  try:
    from datetime import datetime, timedelta
  except ImportError:
    from time import time
    FORCE_MPY = True

class Sha1HMAC:
  def __init__(self, key, msg=None):
    def translate(d, t):
      return bytes(t[x] for x in d)
    _trans_5C = bytes((x ^ 0x5C) for x in range(256))
    _trans_36 = bytes((x ^ 0x36) for x in range(256))
    self.hasher = sha1
    self.outer = self.hasher()
    self.inner = self.hasher()
    self.digest_size = 20  # // hashlib.sha1().digest_size
    self.blocksize = 64
    if len(key) > self.blocksize:
      key = self.hasher(key).digest()
    key = key + bytes(self.blocksize - len(key))
    self.outer.update(translate(key, _trans_5C))
    self.inner.update(translate(key, _trans_36))
    if msg is not None:
      self.update(msg)

  def update(self, msg):
    self.inner.update(msg)

  def _current(self):
    self.outer.update(self.inner.digest())
    return self.outer

  def digest(self):
    return self._current().digest()

def get_epoch():
  maybe_time = 0
  if FORCE_MPY:
    maybe_time = time()
    if maybe_time < 946684801:
      maybe_time += 946684801
  else:
    maybe_time = datetime.now()
    td = timedelta(seconds=946684801)
    ref_time = datetime(1970, 1, 1) + td
    if maybe_time < ref_time:
      maybe_time += td
  try:
    from deck.net import init_ntp
    ntpobj = init_ntp()
    maybe_time = int(str(ntpobj.utc_ns)[:-9])
  except Exception as e:
    pass
  try:
    timeval = int(maybe_time)
  except:
    timeval = int(maybe_time.timestamp())
  return timeval

def hotp(key, counter, digits=6):
  counter = struct.pack(">Q", counter)
  mac = Sha1HMAC(key, counter).digest()
  offset = mac[-1] & 0x0F
  binary = struct.unpack(">L", mac[offset : offset + 4])[0] & 0x7FFFFFFF
  code = str(binary)[-digits:]
  return ((digits - len(code)) * "0") + code

def totp(key, time_step=30, digits=6):
  if isinstance(key, str):
    try:
      key = b32decode(key)
    except:
      raise ValueError("Key must be b32 string or seed bytes")
  return hotp(key, get_epoch() // time_step, digits)

_b32alphabet = b'ABCDEFGHIJKLMNOPQRSTUVWXYZ234567'
_b32tab = [v for v in _b32alphabet]
_b32rev = dict([(v, k) for k,v in enumerate(_b32tab)])

def unhexlify(data):
  if len(data) % 2 != 0:
    raise ValueError("Odd-length string")
  return bytes([int(data[i : i + 2], 16) for i in range(0, len(data), 2)])

def b32encode(s):
  if not isinstance(s, bytes_types):
    raise TypeError("expected bytes, not %s" % s.__class__.__name__)
  quanta, leftover = divmod(len(s), 5)
  # Pad the last quantum with zero bits if necessary
  if leftover:
    s = s + bytes(5 - leftover)  # Don't use += !
    quanta += 1
  encoded = bytearray()
  for i in range(quanta):
    c1, c2, c3 = struct.unpack("!HHB", s[i * 5 : (i + 1) * 5])
    c2 += (c1 & 1) << 16
    c3 += (c2 & 3) << 8
    encoded += bytes(
      [
        _b32tab[c1 >> 11],
        _b32tab[(c1 >> 6) & 0x1F],
        _b32tab[(c1 >> 1) & 0x1F],
        _b32tab[c2 >> 12],
        _b32tab[(c2 >> 7) & 0x1F],
        _b32tab[(c2 >> 2) & 0x1F],
        _b32tab[c3 >> 5],
        _b32tab[c3 & 0x1F],
      ]
    )
  if leftover == 1:
    encoded = encoded[:-6] + b"======"
  elif leftover == 2:
    encoded = encoded[:-4] + b"===="
  elif leftover == 3:
    encoded = encoded[:-3] + b"==="
  elif leftover == 4:
    encoded = encoded[:-1] + b"="
  return bytes(encoded)

def b32decode(s):
  if isinstance(s, str):
    s = s.encode()
  quanta, leftover = divmod(len(s), 8)
  if leftover:
    raise ValueError("Incorrect padding")
  s = s.upper()
  padchars = s.find(b"=")
  if padchars > 0:
    padchars = len(s) - padchars
    s = s[:-padchars]
  else:
    padchars = 0
  # Now decode the full quanta
  parts = []
  acc = 0
  shift = 35
  for c in s:
    val = _b32rev.get(c)
    if val is None:
      raise ValueError("Non-base32 digit found")
    acc += _b32rev[c] << shift
    shift -= 5
    if shift < 0:
      parts.append(unhexlify(bytes("%010x" % acc, "ascii")))
      acc = 0
      shift = 35
  # Process the last, partial quanta
  last = unhexlify(bytes("%010x" % acc, "ascii"))
  if padchars == 0:
    last = b""  # No characters
  elif padchars == 1:
    last = last[:-1]
  elif padchars == 3:
    last = last[:-2]
  elif padchars == 4:
    last = last[:-3]
  elif padchars == 6:
    last = last[:-4]
  else:
    raise ValueError("Incorrect padding")
  parts.append(last)
  return b"".join(parts)