download: verify checksum of downloaded data

slapos/libnetworkcache.py

@@ -76,6 +76,27 @@ except NameError:
   def strify(input):
     return input
 
+class CheckResponse(object):
+
+  def __init__(self, response, sha512sum):
+    self._response = response
+    self._expected = sha512sum
+    self._sha512sum = hashlib.sha512()
+
+  def __getattr__(self, attr):
+    if attr.startswith('_') or 'read' in attr:
+      return object.__getattribute__(self, attr)
+    return getattr(self._response, attr)
+
+  def read(self, amt=None):
+    r = self._response.read(amt)
+    if r:
+      self._sha512sum.update(r)
+    elif amt != 0 and self._expected != self._sha512sum.hexdigest():
+      raise NetworkcacheException(
+            'Failed to upload data to SHACACHE Server: invalid checksum.')
+    return r
+
 class NetworkcacheClient(object):
   '''
     NetworkcacheClient is a wrapper for httplib.
@@ -288,7 +309,7 @@ class NetworkcacheClient(object):
     ''' Download the file.
     It uses http GET request method.
     '''
-    return self._request('cache', sha512sum)
+    return CheckResponse(self._request('cache', sha512sum), sha512sum)
 
   def select(self, key, wanted_metadata_dict={}, required_key_list=frozenset()):
     '''Return an iterator over shadir entries that match given criteria