浏览代码

Reworked packer format again

Jonas Borgström 14 年之前
父节点
当前提交
ecc8b1e7c1
共有 5 个文件被更改,包括 48 次插入40 次删除
  1. 7 0
      darc/__init__.py
  2. 15 9
      darc/archive.py
  3. 8 7
      darc/cache.py
  4. 16 22
      darc/keychain.py
  5. 2 2
      darc/store.py

+ 7 - 0
darc/__init__.py

@@ -5,3 +5,10 @@ NS_ARCHIVE_METADATA = 1
 NS_ARCHIVE_CHUNKS = 2
 NS_ARCHIVE_ITEMS = 3
 
+PACKET_ENCRYPT_READ   = 2 ** 7
+PACKET_ENCRYPT_CREATE = 2 ** 6
+PACKET_CHUNK            = 1 | PACKET_ENCRYPT_READ
+PACKET_ARCHIVE_METADATA = 2 | PACKET_ENCRYPT_READ
+PACKET_ARCHIVE_ITEMS    = 3 | PACKET_ENCRYPT_READ
+PACKET_ARCHIVE_CHUNKS   = 1 | PACKET_ENCRYPT_CREATE
+

+ 15 - 9
darc/archive.py

@@ -7,7 +7,8 @@ import stat
 import sys
 from xattr import xattr, XATTR_NOFOLLOW
 
-from . import NS_ARCHIVE_METADATA, NS_ARCHIVE_ITEMS, NS_ARCHIVE_CHUNKS, NS_CHUNK
+from . import NS_ARCHIVE_METADATA, NS_ARCHIVE_ITEMS, NS_ARCHIVE_CHUNKS, NS_CHUNK, \
+    PACKET_ARCHIVE_METADATA, PACKET_ARCHIVE_ITEMS, PACKET_ARCHIVE_CHUNKS, PACKET_CHUNK
 from ._speedups import chunkify
 from .helpers import uid2user, user2uid, gid2group, group2gid, IntegrityError
 
@@ -35,15 +36,17 @@ class Archive(object):
     def load(self, id):
         self.id = id
         try:
-            data, self.hash = self.keychain.decrypt(self.store.get(NS_ARCHIVE_METADATA, self.id))
+            kind, data, self.hash = self.keychain.decrypt(self.store.get(NS_ARCHIVE_METADATA, self.id))
         except self.store.DoesNotExist:
             raise self.DoesNotExist
+        assert kind == PACKET_ARCHIVE_METADATA
         self.metadata = msgpack.unpackb(data)
         assert self.metadata['version'] == 1
 
     def get_chunks(self):
         for id in self.metadata['chunks_ids']:
-            data, hash = self.keychain.decrypt(self.store.get(NS_ARCHIVE_CHUNKS, id))
+            magic, data, hash = self.keychain.decrypt(self.store.get(NS_ARCHIVE_CHUNKS, id))
+            assert magic == PACKET_ARCHIVE_CHUNKS
             assert hash == id
             chunks = msgpack.unpackb(data)
             for chunk in chunks:
@@ -51,7 +54,8 @@ class Archive(object):
 
     def get_items(self):
         for id in self.metadata['items_ids']:
-            data, items_hash = self.keychain.decrypt(self.store.get(NS_ARCHIVE_ITEMS, id))
+            magic, data, items_hash = self.keychain.decrypt(self.store.get(NS_ARCHIVE_ITEMS, id))
+            assert magic == PACKET_ARCHIVE_ITEMS
             assert items_hash == id
             items = msgpack.unpackb(data)
             for item in items:
@@ -63,7 +67,7 @@ class Archive(object):
             self.flush_items()
 
     def flush_items(self):
-        data, hash = self.keychain.encrypt_read(msgpack.packb(self.items))
+        data, hash = self.keychain.encrypt(PACKET_ARCHIVE_ITEMS, msgpack.packb(self.items))
         self.store.put(NS_ARCHIVE_ITEMS, hash, data)
         self.items_ids.append(hash)
         self.items = []
@@ -72,7 +76,7 @@ class Archive(object):
         chunks = []
         ids = []
         def flush(chunks):
-            data, hash = self.keychain.encrypt_create(msgpack.packb(chunks))
+            data, hash = self.keychain.encrypt(PACKET_ARCHIVE_CHUNKS, msgpack.packb(chunks))
             self.store.put(NS_ARCHIVE_CHUNKS, hash, data)
             ids.append(hash)
         for id, (count, size) in cache.chunk_counts.iteritems():
@@ -98,7 +102,7 @@ class Archive(object):
             'username': getuser(),
             'time': datetime.utcnow().isoformat(),
         }
-        data, self.hash = self.keychain.encrypt_read(msgpack.packb(metadata))
+        data, self.hash = self.keychain.encrypt(PACKET_ARCHIVE_METADATA, msgpack.packb(metadata))
         self.store.put(NS_ARCHIVE_METADATA, self.id, data)
         self.store.commit()
 
@@ -149,7 +153,8 @@ class Archive(object):
                 with open(path, 'wb') as fd:
                     for id in item['chunks']:
                         try:
-                            data, hash = self.keychain.decrypt(self.store.get(NS_CHUNK, id))
+                            magic, data, hash = self.keychain.decrypt(self.store.get(NS_CHUNK, id))
+                            assert magic == PACKET_CHUNK
                             if self.keychain.id_hash(data) != id:
                                 raise IntegrityError('chunk hash did not match')
                             fd.write(data)
@@ -185,7 +190,8 @@ class Archive(object):
     def verify_file(self, item):
         for id in item['chunks']:
             try:
-                data, hash = self.keychain.decrypt(self.store.get(NS_CHUNK, id))
+                magic, data, hash = self.keychain.decrypt(self.store.get(NS_CHUNK, id))
+                assert magic == PACKET_CHUNK
                 if self.keychain.id_hash(data) != id:
                     raise IntegrityError('chunk id did not match')
             except IntegrityError:

+ 8 - 7
darc/cache.py

@@ -1,7 +1,7 @@
 import msgpack
 import os
 
-from . import NS_ARCHIVE_CHUNKS, NS_CHUNK
+from . import NS_ARCHIVE_CHUNKS, NS_CHUNK, PACKET_ARCHIVE_CHUNKS, PACKET_CHUNK
 
 
 class Cache(object):
@@ -27,8 +27,8 @@ class Cache(object):
         if not os.path.exists(self.path):
             return
         with open(self.path, 'rb') as fd:
-            data, hash = self.keychain.decrypt(fd.read())
-            cache = msgpack.unpackb(data)
+            #data, hash = self.keychain.decrypt(fd.read())
+            cache = msgpack.unpackb(fd.read())
         assert cache['version'] == 1
         self.chunk_counts = cache['chunk_counts']
         self.file_chunks = cache['file_chunks']
@@ -44,7 +44,8 @@ class Cache(object):
             if len(id) != 32:
                 import ipdb
                 ipdb.set_trace()
-            data, hash = self.keychain.decrypt(self.store.get(NS_ARCHIVE_CHUNKS, id))
+            magic, data, hash = self.keychain.decrypt(self.store.get(NS_ARCHIVE_CHUNKS, id))
+            assert magic == PACKET_ARCHIVE_CHUNKS
             chunks = msgpack.unpackb(data)
             for id, size in chunks:
                 try:
@@ -69,17 +70,17 @@ class Cache(object):
                 'chunk_counts': self.chunk_counts,
                 'file_chunks': dict(self.filter_file_chunks()),
         }
-        data, hash = self.keychain.encrypt_create(msgpack.packb(cache))
+#        data, hash = self.keychain.encrypt_create(msgpack.packb(cache))
         cachedir = os.path.dirname(self.path)
         if not os.path.exists(cachedir):
             os.makedirs(cachedir)
         with open(self.path, 'wb') as fd:
-            fd.write(data)
+            fd.write(msgpack.packb(cache))
 
     def add_chunk(self, id, data):
         if self.seen_chunk(id):
             return self.chunk_incref(id)
-        data, hash = self.keychain.encrypt_read(data)
+        data, hash = self.keychain.encrypt(PACKET_CHUNK, data)
         csize = len(data)
         self.store.put(NS_CHUNK, id, data)
         self.chunk_counts[id] = (1000001, csize)

+ 16 - 22
darc/keychain.py

@@ -11,6 +11,7 @@ from Crypto.PublicKey import RSA
 from Crypto.Util import Counter
 from Crypto.Util.number import bytes_to_long, long_to_bytes
 
+from . import PACKET_ENCRYPT_READ, PACKET_ENCRYPT_CREATE
 from .helpers import IntegrityError, zero_pad
 from .oaep import OAEP
 
@@ -140,28 +141,21 @@ class Keychain(object):
         """
         return HMAC.new(self.aes_id, data, SHA256).digest()
 
-    def _encrypt(self, id, rsa_key, key, data):
+    def encrypt(self, magic, data):
         """Helper function used by `encrypt_read` and `encrypt_create`
         """
         data = zlib.compress(data)
         nonce = long_to_bytes(self.counter.next_value(), 8)
-        data = nonce + rsa_key + AES.new(key, AES.MODE_CTR, '', counter=self.counter).encrypt(data)
+        if magic & PACKET_ENCRYPT_READ:
+            data = ''.join((nonce, self.read_encrypted,
+                            AES.new(self.read_key, AES.MODE_CTR, '',
+                                    counter=self.counter).encrypt(data)))
+        elif magic & PACKET_ENCRYPT_CREATE:
+            data = ''.join((nonce, self.create_encrypted,
+                            AES.new(self.create_key, AES.MODE_CTR, '',
+                                    counter=self.counter).encrypt(data)))
         hash = self.id_hash(data)
-        return ''.join((id, hash, data)), hash
-
-    def encrypt_read(self, data):
-        """Encrypt `data` using the AES "read" key
-
-        An RSA encrypted version of the AES key is included in the header
-        """
-        return self._encrypt(self.READ, self.read_encrypted, self.read_key, data)
-
-    def encrypt_create(self, data, iv=None):
-        """Encrypt `data` using the AES "create" key
-
-        An RSA encrypted version of the AES key is included in the header
-        """
-        return self._encrypt(self.CREATE, self.create_encrypted, self.create_key, data)
+        return ''.join((chr(magic), hash, data)), hash
 
     def _decrypt_key(self, data, rsa_key):
         """Helper function used by `decrypt`
@@ -175,20 +169,20 @@ class Keychain(object):
     def decrypt(self, data):
         """Decrypt `data` previously encrypted by `encrypt_create` or `encrypt_read`
         """
-        type = data[0]
+        magic = ord(data[0])
         hash = data[1:33]
         if self.id_hash(data[33:]) != hash:
             raise IntegrityError('Encryption integrity error')
         nonce = bytes_to_long(data[33:41])
         counter = Counter.new(64, prefix='\0' * 8, initial_value=nonce)
-        if type == self.READ:
+        if magic & PACKET_ENCRYPT_READ:
             key = self._decrypt_key(data[41:297], self.rsa_read)
-        elif type == self.CREATE:
+        elif magic & PACKET_ENCRYPT_CREATE:
             key = self._decrypt_key(data[41:297], self.rsa_create)
         else:
-            raise Exception('Unknown pack type %d found' % ord(type))
+            raise Exception('Unknown pack magic %d found' % magic)
         data = AES.new(key, AES.MODE_CTR, counter=counter).decrypt(data[297:])
-        return zlib.decompress(data), hash
+        return magic, zlib.decompress(data), hash
 
 
 

+ 2 - 2
darc/store.py

@@ -66,7 +66,7 @@ class Store(object):
         self.config.read(os.path.join(path, 'config'))
         if self.config.getint('store', 'version') != 1:
             raise Exception('%s Does not look like a darc store')
-        self.id = self.config.get('store', 'id')
+        self.id = self.config.get('store', 'id').decode('hex')
         self.tid = self.config.getint('state', 'tid')
         next_band = self.config.getint('state', 'next_band')
         max_band_size = self.config.getint('store', 'max_band_size')
@@ -287,7 +287,6 @@ class HashIndex(DictMixin):
 
     def resize(self, capacity=0):
         capacity = capacity or self.buckets.size * 2
-        print 'resizing to', capacity
         if capacity < self.num_entries:
             raise ValueError('HashIndex full')
         new = HashIndex.create(self.path + '.tmp', capacity)
@@ -341,6 +340,7 @@ class BandIO(object):
         fd.seek(offset)
         data = fd.read(self.header_fmt.size)
         size, magic, ns, id = self.header_fmt.unpack(data)
+        assert magic == 0
         return fd.read(size - self.header_fmt.size)
 
     def iter_objects(self, band):