From 9a7b948fe69ade6de49cfa1e357d6cbc40688f43 Mon Sep 17 00:00:00 2001
From: Mark Wooding <mdw@distorted.org.uk>
Date: Sat, 24 May 2014 14:00:03 +0100
Subject: [PATCH] catacomb/pwsafe.py: Eliminate the Buffer class and struct
 module.

All is done using Catacomb's ReadBuffer and WriteBuffer classes.
---
 catacomb/pwsafe.py | 130 +++++++++++++++--------------------------------------
 1 file changed, 36 insertions(+), 94 deletions(-)

diff --git a/catacomb/pwsafe.py b/catacomb/pwsafe.py
index d0d1a35..3af69d4 100644
--- a/catacomb/pwsafe.py
+++ b/catacomb/pwsafe.py
@@ -28,73 +28,6 @@
 
 import catacomb as _C
 import gdbm as _G
-import struct as _S
-
-###--------------------------------------------------------------------------
-### Utilities.
-
-class Buffer (object):
-  """
-  I am a simple gadget for parsing binary strings.
-
-  You should use Catacomb's ReadBuffer instead.
-  """
-
-  def __init__(me, s):
-    """
-    Initialize the buffer with a string S.
-    """
-    me.str = s
-    me.i = 0
-
-  def get(me, n):
-    """
-    Fetch and return the next N bytes from the buffer.
-    """
-    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):
-    """
-    Fetch and return (as a small integer) the next byte from the buffer.
-    """
-    return ord(me.get(1))
-
-  def unpack(me, fmt):
-    """
-    Unpack a structure described by FMT from the next bytes of the buffer.
-
-    Return a tuple containing the unpacked items.
-    """
-    return _S.unpack(fmt, me.get(_S.calcsize(fmt)))
-
-  def getstring(me):
-    """
-    Fetch and return a counted string from the buffer.
-
-    The string is expected to be preceded by its 16-bit length, in network
-    byte order.
-    """
-    return me.get(me.unpack('>H')[0])
-
-  def checkend(me):
-    """
-    Raise an error if the buffer has not been completely consumed.
-    """
-    if me.i != len(me.str):
-      raise ValueError, 'junk at end of buffer'
-
-def _wrapstr(s):
-  """
-  Prefix the string S with its 16-bit length.
-
-  It can be read using Buffer.getstring.  You should use Catacomb's
-  WriteBuffer.putblk16() function instead.
-  """
-  return _S.pack('>H', len(s)) + s
 
 ###--------------------------------------------------------------------------
 ### Underlying cryptography.
@@ -137,27 +70,35 @@ class Crypto (object):
     """
     Encrypt the message PT and return the resulting ciphertext.
     """
-    if me.c.__class__.blksz:
-      iv = _C.rand.block(me.c.__class__.blksz)
+    blksz = me.c.__class__.blksz
+    b = _C.WriteBuffer()
+    if blksz:
+      iv = _C.rand.block(blksz)
       me.c.setiv(iv)
-    else:
-      iv = ''
-    y = iv + me.c.encrypt(pt)
-    t = me.m().hash(y).done()
-    return t + y
+      b.put(iv)
+    b.put(me.c.encrypt(pt))
+    t = me.m().hash(b).done()
+    return t + str(buffer(b))
+
   def decrypt(me, ct):
     """
     Decrypt the ciphertext CT, returning the plaintext.
 
     Raises DecryptError if anything goes wrong.
     """
-    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:])
+    blksz = me.c.__class__.blksz
+    tagsz = me.m.__class__.tagsz
+    b = _C.ReadBuffer(ct)
+    t = b.get(tagsz)
+    h = me.m()
+    if blksz:
+      iv = b.get(blksz)
+      me.c.setiv(iv)
+      h.hash(iv)
+    x = b.get(b.left)
+    h.hash(x)
+    if t != h.done(): raise DecryptError
+    return me.c.decrypt(x)
 
 class PPK (Crypto):
   """
@@ -271,13 +212,13 @@ class PW (object):
     tag = me.db['tag']
     ppk = PPK(_C.ppread(tag), c, h, m, me.db['salt'])
     try:
-      buf = Buffer(ppk.decrypt(me.db['key']))
+      b = _C.ReadBuffer(ppk.decrypt(me.db['key']))
     except DecryptError:
       _C.ppcancel(tag)
       raise
-    me.ck = buf.getstring()
-    me.mk = buf.getstring()
-    buf.checkend()
+    me.ck = b.getblk16()
+    me.mk = b.getblk16()
+    if not b.endp: raise ValueError, 'trailing junk'
 
     ## Set the key, and stash it and the tag-hashing secret.
     me.k = Crypto(c, h, m, me.ck, me.mk)
@@ -309,7 +250,7 @@ class PW (object):
     db['cipher'] = c.name
     db['hash'] = h.name
     db['mac'] = m.name
-    db['key'] = ppk.encrypt(_wrapstr(ck) + _wrapstr(mk))
+    db['key'] = ppk.encrypt(_C.WriteBuffer().putblk16(ck).putblk16(mk))
     db['magic'] = k.encrypt(_C.rand.block(h.hashsz))
 
   def keyxform(me, key):
@@ -329,17 +270,18 @@ class PW (object):
     _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['key'] = \
+        ppk.encrypt(_C.WriteBuffer().putblk16(me.ck).putblk16(me.mk))
     me.db['salt'] = ppk.salt
 
   def pack(me, key, value):
     """
     Pack the KEY and VALUE into a ciphertext, and return it.
     """
-    w = _wrapstr(key) + _wrapstr(value)
-    pl = (len(w) + 255) & ~255
-    w += '\0' * (pl - len(w))
-    return me.k.encrypt(w)
+    b = _C.WriteBuffer()
+    b.putblk16(key).putblk16(value)
+    b.zero(((b.size + 255) & ~255) - b.size)
+    return me.k.encrypt(b)
 
   def unpack(me, ct):
     """
@@ -347,9 +289,9 @@ class PW (object):
 
     Might raise DecryptError, of course.
     """
-    buf = Buffer(me.k.decrypt(ct))
-    key = buf.getstring()
-    value = buf.getstring()
+    b = _C.ReadBuffer(me.k.decrypt(ct))
+    key = b.getblk16()
+    value = b.getblk16()
     return key, value
 
   ## Mapping protocol.
-- 
2.11.0