Various other little bits.

[catacomb-python] / catacomb / pwsafe.py

# -*-python-*-

import catacomb as _C
import gdbm as _G
import struct as _S

class DecryptError (Exception):
  pass

class Crypto (object):
  def __init__(me, c, h, m, ck, mk):
    me.c = c(ck)
    me.m = m(mk)
    me.h = h
  def encrypt(me, pt):
    if me.c.__class__.blksz:
      iv = _C.rand.block(me.c.__class__.blksz)
      me.c.setiv(iv)
    else:
      iv = ''
    y = iv + me.c.encrypt(pt)
    t = me.m().hash(y).done()
    return t + y
  def decrypt(me, ct):
    t = ct[:me.m.__class__.tagsz]
    y = ct[me.m.__class__.tagsz:]
    if t != me.m().hash(y).done():
      raise DecryptError
    iv = y[:me.c.__class__.blksz]
    if me.c.__class__.blksz: me.c.setiv(iv)
    return me.c.decrypt(y[me.c.__class__.blksz:])
  
class PPK (Crypto):
  def __init__(me, pp, c, h, m, salt = None):
    if not salt: salt = _C.rand.block(h.hashsz)
    tag = '%s\0%s' % (pp, salt)
    Crypto.__init__(me, c, h, m,
                  h().hash('cipher:' + tag).done(),
                  h().hash('mac:' + tag).done())
    me.salt = salt

class Buffer (object):
  def __init__(me, s):
    me.str = s
    me.i = 0
  def get(me, n):
    i = me.i
    if n + i > len(me.str):
      raise IndexError, 'buffer underflow'
    me.i += n
    return me.str[i:i + n]
  def getbyte(me):
    return ord(me.get(1))
  def unpack(me, fmt):
    return _S.unpack(fmt, me.get(_S.calcsize(fmt)))
  def getstring(me):
    return me.get(me.unpack('>H')[0])
  def checkend(me):
    if me.i != len(me.str):
      raise ValueError, 'junk at end of buffer'

def _wrapstr(s):
  return _S.pack('>H', len(s)) + s

class PWIter (object):
  def __init__(me, pw):
    me.pw = pw
    me.k = me.pw.db.firstkey()
  def next(me):
    k = me.k
    while True:
      if k is None:
        raise StopIteration
      if k[0] == '$':
        break
      k = me.pw.db.nextkey(k)
    me.k = me.pw.db.nextkey(k)
    return me.pw.unpack(me.pw.db[k])[0]
class PW (object):
  def __init__(me, file, mode = 'r'):
    me.db = _G.open(file, mode)
    c = _C.gcciphers[me.db['cipher']]
    h = _C.gchashes[me.db['hash']]
    m = _C.gcmacs[me.db['mac']]
    tag = me.db['tag']
    ppk = PPK(_C.ppread(tag), c, h, m, me.db['salt'])
    try:
      buf = Buffer(ppk.decrypt(me.db['key']))
    except DecryptError:
      _C.ppcancel(tag)
      raise
    me.ck = buf.getstring()
    me.mk = buf.getstring()
    buf.checkend()
    me.k = Crypto(c, h, m, me.ck, me.mk)
    me.magic = me.k.decrypt(me.db['magic'])
  def keyxform(me, key):
    return '$' + me.k.h().hash(me.magic).hash(key).done()
  def changepp(me):
    tag = me.db['tag']
    _C.ppcancel(tag)
    ppk = PPK(_C.ppread(tag, _C.PMODE_VERIFY),
              me.k.c.__class__, me.k.h, me.k.m.__class__)
    me.db['key'] = ppk.encrypt(_wrapstr(me.ck) + _wrapstr(me.mk))
    me.db['salt'] = ppk.salt
  def pack(me, key, value):
    w = _wrapstr(key) + _wrapstr(value)
    pl = (len(w) + 255) & ~255
    w += '\0' * (pl - len(w))
    return me.k.encrypt(w)
  def unpack(me, p):
    buf = Buffer(me.k.decrypt(p))
    key = buf.getstring()
    value = buf.getstring()
    return key, value
  def __getitem__(me, key):
    try:
      return me.unpack(me.db[me.keyxform(key)])[1]
    except KeyError:
      raise KeyError, key
  def __setitem__(me, key, value):
    me.db[me.keyxform(key)] = me.pack(key, value)
  def __delitem__(me, key):
    try:
      del me.db[me.keyxform(key)]
    except KeyError:
      raise KeyError, key
  def __iter__(me):
    return PWIter(me)