X-Git-Url: https://git.distorted.org.uk/~mdw/catacomb-python/blobdiff_plain/b35fdbe6fd7d9d46d5e195ae8cfd14d5729d3234..6bd22b53a63463d18ee6b1bd4d0ca0f02a8f5b99:/catacomb/__init__.py

diff --git a/catacomb/__init__.py b/catacomb/__init__.py
index d4eac30..7ea8d57 100644
--- a/catacomb/__init__.py
+++ b/catacomb/__init__.py
@@ -23,10 +23,14 @@
 ### along with Catacomb/Python; if not, write to the Free Software Foundation,
 ### Inc., 59 Temple Place - Suite 330, Boston, MA 02111-1307, USA.
 
+from __future__ import with_statement
+
 import _base
 import types as _types
 from binascii import hexlify as _hexify, unhexlify as _unhexify
+from contextlib import contextmanager as _ctxmgr
 from sys import argv as _argv
+from struct import pack as _pack
 
 ###--------------------------------------------------------------------------
 ### Basic stuff.
@@ -160,6 +164,96 @@ class _tmp:
 _augment(GHash, _tmp)
 _augment(Poly1305Hash, _tmp)
 
+class _HashBase (object):
+  ## The standard hash methods.  Assume that `hash' is defined and returns
+  ## the receiver.
+  def hashu8(me, n): return me.hash(_pack('B', n))
+  def hashu16l(me, n): return me.hash(_pack('<H', n))
+  def hashu16b(me, n): return me.hash(_pack('>H', n))
+  hashu16 = hashu16b
+  def hashu32l(me, n): return me.hash(_pack('<L', n))
+  def hashu32b(me, n): return me.hash(_pack('>L', n))
+  hashu32 = hashu32b
+  def hashu64l(me, n): return me.hash(_pack('<Q', n))
+  def hashu64b(me, n): return me.hash(_pack('>Q', n))
+  hashu64 = hashu64b
+  def hashbuf8(me, s): return me.hashu8(len(s)).hash(s)
+  def hashbuf16l(me, s): return me.hashu16l(len(s)).hash(s)
+  def hashbuf16b(me, s): return me.hashu16b(len(s)).hash(s)
+  hashbuf16 = hashbuf16b
+  def hashbuf32l(me, s): return me.hashu32l(len(s)).hash(s)
+  def hashbuf32b(me, s): return me.hashu32b(len(s)).hash(s)
+  hashbuf32 = hashbuf32b
+  def hashbuf64l(me, s): return me.hashu64l(len(s)).hash(s)
+  def hashbuf64b(me, s): return me.hashu64b(len(s)).hash(s)
+  hashbuf64 = hashbuf64b
+  def hashstrz(me, s): return me.hash(s).hashu8(0)
+
+class _ShakeBase (_HashBase):
+
+  ## Python gets really confused if I try to augment `__new__' on native
+  ## classes, so wrap and delegate.  Sorry.
+  def __init__(me, perso = '', *args, **kw):
+    super(_ShakeBase, me).__init__(*args, **kw)
+    me._h = me._SHAKE(perso = perso, func = me._FUNC)
+
+  ## Delegate methods...
+  def copy(me): new = me.__class__(); new._copy(me)
+  def _copy(me, other): me._h = other._h
+  def hash(me, m): me._h.hash(m); return me
+  def xof(me): me._h.xof(); return me
+  def get(me, n): return me._h.get(n)
+  def mask(me, m): return me._h.mask(m)
+  def done(me, n): return me._h.done(n)
+  def check(me, h): return ctstreq(h, me.done(len(h)))
+  @property
+  def state(me): return me._h.state
+  @property
+  def buffered(me): return me._h.buffered
+  @property
+  def rate(me): return me._h.rate
+
+class _tmp:
+  def check(me, h):
+    return ctstreq(h, me.done(len(h)))
+  def leftenc(me, n):
+    nn = MP(n).storeb()
+    return me.hashu8(len(nn)).hash(nn)
+  def rightenc(me, n):
+    nn = MP(n).storeb()
+    return me.hash(nn).hashu8(len(nn))
+  def stringenc(me, str):
+    return me.leftenc(8*len(str)).hash(str)
+  def bytepad_before(me):
+    return me.leftenc(me.rate)
+  def bytepad_after(me):
+    if me.buffered: me.hash(me._Z[:me.rate - me.buffered])
+    return me
+  @_ctxmgr
+  def bytepad(me):
+    me.bytepad_before()
+    yield me
+    me.bytepad_after()
+_augment(Shake, _tmp)
+_augment(_ShakeBase, _tmp)
+Shake._Z = _ShakeBase._Z = ByteString(200*'\0')
+
+class KMAC (_ShakeBase):
+  _FUNC = 'KMAC'
+  def __init__(me, k, *arg, **kw):
+    super(KMAC, me).__init__(*arg, **kw)
+    with me.bytepad(): me.stringenc(k)
+  def done(me, n = -1):
+    if n < 0: n = me._TAGSZ
+    me.rightenc(8*n)
+    return super(KMAC, me).done(n)
+  def xof(me):
+    me.rightenc(0)
+    return super(KMAC, me).xof()
+
+class KMAC128 (KMAC): _SHAKE = Shake128; _TAGSZ = 16
+class KMAC256 (KMAC): _SHAKE = Shake256; _TAGSZ = 32
+
 ###--------------------------------------------------------------------------
 ### NaCl `secretbox'.