code.grnet.gr Git - pithos/blob - tools/psend

   1 #!/usr/bin/env python
   2
   3 import os
   4 import hashlib
   5 import sys
   6
   7 from binascii import hexlify, unhexlify
   8 from cStringIO import StringIO
   9
  10 from lib.client import Pithos_Client, Fault
  11 from lib.util import get_user, get_auth, get_server, get_api
  12
  13
  14 # XXX Get these from container...
  15 BLOCK_SIZE = 4 * 1024 * 1024
  16 BLOCK_HASH = 'sha256'
  17
  18
  19 def file_read_iterator(fp, size=1024):
  20     while True:
  21         data = fp.read(size)
  22         if not data:
  23             break
  24         yield data
  25
  26
  27 class HashMap(list):
  28
  29     def __init__(self, f):
  30         super(HashMap, self).__init__()
  31         self.load(f)
  32
  33     def _hash_raw(self, v):
  34         h = hashlib.new(BLOCK_HASH)
  35         h.update(v)
  36         return h.digest()
  37
  38     def _hash_block(self, v):
  39         return self._hash_raw(v.rstrip('\x00'))
  40
  41     def hash(self):
  42         if len(self) == 0:
  43             return self._hash_raw('')
  44         if len(self) == 1:
  45             return self.__getitem__(0)
  46
  47         h = list(self)
  48         s = 2
  49         while s < len(h):
  50             s = s * 2
  51         h += [('\x00' * len(h[0]))] * (s - len(h))
  52         while len(h) > 1:
  53             h = [self._hash_raw(h[x] + h[x + 1]) for x in range(0, len(h), 2)]
  54         return h[0]
  55
  56     def load(self, f):
  57         with open(f) as fp:
  58             for block in file_read_iterator(fp, BLOCK_SIZE):
  59                 self.append(self._hash_block(block))
  60
  61
  62 def smart_upload(client, file):
  63     dest_container = 'pithos'
  64     dest_object = os.path.split(file)[-1]
  65
  66     size = os.path.getsize(file)
  67     hashes = HashMap(sys.argv[1])
  68     map = {'bytes': size, 'hashes': [hexlify(x) for x in hashes]}
  69
  70     try:
  71         client.create_object_by_hashmap(dest_container, dest_object, map)
  72     except Fault, fault:
  73         if fault.status != 409:
  74             raise
  75     else:
  76         return
  77
  78     missing = fault.data.split('\n')
  79     if '' in missing:
  80         del missing[missing.index(''):]
  81
  82     with open(file) as fp:
  83         for hash in missing:
  84             offset = hashes.index(unhexlify(hash)) * BLOCK_SIZE
  85             fp.seek(offset)
  86             block = fp.read(BLOCK_SIZE)
  87             client.create_object('pithos', '.upload', StringIO(block))
  88
  89     client.create_object_by_hashmap(dest_container, dest_object, map)
  90
  91
  92 if __name__ == '__main__':
  93     if len(sys.argv) != 2 or not os.path.isfile(sys.argv[1]):
  94         print 'syntax: %s <file>' % sys.argv[0]
  95         sys.exit(1)
  96
  97     client = Pithos_Client(get_server(), get_auth(), get_user())
  98     smart_upload(client, sys.argv[1])