1 files changed, 68 insertions, 21 deletions
diff --git a/crypto/ocb.py b/crypto/ocb.py
index 08814c7..7383ee6 100644
--- a/crypto/ocb.py
+++ b/crypto/ocb.py
@@ -2,27 +2,74 @@
 from aespython import key_expander, aes_cipher, ocb_mode
 import json
 import base64
-def hex2byte(hexstring):
-    return map(ord, hexstring.decode("hex"))
-def byte2hex(byteslist):
-    return str(bytearray(byteslist)).encode("hex")
 
+class sjcl(ocb_mode.OCBMode):
+    """convenience class to wrap the sjcl objects encryption and decryption"""
+    def __init__(self, key, keylen):
+        KE = key_expander.KeyExpander(keylen)
+        expended_key = KE.expand(key)
+        cipher = aes_cipher.AESCipher(expanded_key)
+        ocb_mode.OCBMode.__init__(self, cipher, len(key))
+
+    @staticmethod
+    def hex2byte(hexstring):
+        """convert a hexstring to a list of bytes"""
+        return map(ord, hexstring.decode("hex"))
+
+    @staticmethod
+    def byte2hex(byteslist):
+        """convert a list of bytes to a hexstring"""
+        return str(bytearray(byteslist)).encode("hex")
+    
+    @staticmethod
+    def pad(b64string):
+        """pad a base64 string with = so that its lengh is divisible by 4"""
+        return b64string + ''.join(['='] * (len(b64string) % 4))
+
+    def encrypt_sjclobject(self, plaintext, header):
+        """encrypt a message and return in in the sjcl object format"""
+        tag, ciphertext = ocb_mode.OCBMode.encrypt_block(self, plaintext, header)
+        sjcl_object = {"iv": base64.b64encode(bytearray(self._iv)).rstrip("="),
+                       "mode": "ocb2",
+                       "ts": 128,
+                       "ct": base64.b64encode(bytearray(ciphertext) +
+                                              bytearray(tag)).rstrip("="),
+                       "adata": base64.b64encode(header).rstrip("=")
+                  }
+        return json.dumps(sjcl_object)
+
+    def decrypt_sjclobject(self, objectstring):
+        """takes a sjcl object and return a list of bytes."""
+        sjcl_object = json.loads(objectstring)
+        iv = map(ord, base64.b64decode(pad(sjcl_object['iv'])))
+        self.set_iv(iv)
+        ct = map(ord, base64.b64decode(pad(sjcl_object['ct'])))
+        try:
+            header = map(ord, base64.b64decode(pad(sjcl_object['adata'])))
+        except KeyError:
+            header = []
+        #sjcl appends the tag at the end of the cipher text (it's 16 bytes long)
+        tag = ct[-16:]
+        ciphertext = ct[:-16]
+        isValid, plaintext = ocb_mode.OCBMode.decrypt_block(self, header, ciphertext, tag)
+        if not isValid:
+            raise Exception("can't authenticate the data")
+        else:
+            return plaintext
+        
 if __name__=="__main__":
-    KE = key_expander.KeyExpander(128)
-    key = hex2byte('12538243c49f1c58e6f7b0687bbd65b2')
-    expanded_key = KE.expand(key)
-    aes_cipher_128 = aes_cipher.AESCipher(expanded_key)
-    aes_ocb_128 = ocb_mode.OCBMode(aes_cipher_128, 16)
-    iv = hex2byte('250c3041c00a605a4100e264abbc588b')
-    aes_ocb_128.set_iv(iv)
-    plaintext = "La chaire est triste, hélas ! et j'ai lu tous les livres.".encode("hex")
-    header = "".encode("hex")
-    (tag,ciphertext) = aes_ocb_128.encrypt_block(hex2byte(plaintext), hex2byte(header))
-    sjcl_object = {"iv": base64.b64encode(bytearray(iv)).rstrip("="),
-                 "mode": "ocb2",
-                 "tag": 128,
-                 "ct": base64.b64encode(bytearray(ciphertext) +
-                                        bytearray(tag)).rstrip("=")
-                 }
-    print json.dumps(sjcl_object)
+    key = sjcl.hex2byte('12538243c49f1c58e6f7b0687bbd65b2')
+    iv = sjcl.hex2byte('250c3041c00a605a4100e264abbc588b')
+    sjclmode = sjcl(key, 128)
+    sjclmode.set_iv(iv)
+    plaintext = map(ord, "La chaire est triste, hélas ! et j'ai lu \
+tous les livres.")
+    header = bytearray("Mallarmé")
+    test = sjclmode.encrypt_sjclobject(plaintext, header)   
+    print test
+    print str(bytearray(sjclmode.decrypt_sjclobject(test)))
+
+
+
+