database_id.py

import os
import time

from construct import (
  BitStruct, Default, ExprAdapter, Flag, GreedyBytes, Int16ul, Padding, Struct,
  Transformed, obj_,
)
from telethon.utils import _decode_telegram_base64, _encode_telegram_base64

from util import tt_from_hex, tt_to_hex


class UUID:
  Flags = Transformed(
    BitStruct(
      Padding(4),
      'is_major' / Default(Flag, False),
      'skip_saving' / Default(Flag, False),
      'timestamp' / Default(Flag, True),
      '_unused' / Default(Flag, False),
    ),
    lambda v: v.translate(tt_from_hex), 1,
    lambda v: v.translate(tt_to_hex), 1
  )
  DBKey = Struct(
    'timestamp' / Default(
      ExprAdapter(Int16ul, obj_ << 16, (obj_ >> 16) & 0xFFFF),
      lambda this: int(time.time())
    ),
    'random' / Default(GreedyBytes, lambda this: os.urandom(45)),
  )

  def __init__(self, old=''):
    flag_byte = old[0].encode('utf8') if old else UUID.Flags.build({})
    self.flags = UUID.Flags.parse(flag_byte)
    if old:
      self.db_key = old[1:]
    else:
      self.db_key = _encode_telegram_base64(UUID.DBKey.build({}))

  @property
  def is_major(self):
    return self.flags.is_major

  def read_timestamp(self):
    if not self.flags.timestamp:
      return
    parsed = UUID.DBKey.parse(_decode_telegram_base64(self.db_key))
    return parsed.timestamp

  def get_str(self, is_major=None, skip_saving=None):
    if is_major is not None:
      self.flags.is_major = is_major
    if skip_saving is not None:
      self.flags.skip_saving = skip_saving
    flag = UUID.Flags.build(self.flags)
    return flag.decode('ascii') + self.db_key