Adding support to signature...

slapos/libnetworkcache.py

@@ -21,6 +21,10 @@ import os
 import tempfile
 import urllib
 import urlparse
+import M2Crypto
+
+
+_MARKER = (None, '')
 
 
 class NetworkcacheClient(object):
@@ -48,7 +52,8 @@ class NetworkcacheClient(object):
     return_dict['port'] = parsed_url.port
     return return_dict
 
-  def __init__(self, shacache, shadir):
+  def __init__(self, shacache, shadir,
+               signature_private_file=None, signature_public_file=None):
     ''' Set the initial values. '''
     # ShaCache Properties
     for k, v in self.parseUrl(shacache).iteritems():
@@ -59,6 +64,9 @@ class NetworkcacheClient(object):
     for k, v in self.parseUrl(shadir).iteritems():
       setattr(self, 'shadir_%s' % k, v)
 
+    self.signature_private_file = signature_private_file
+    self.signature_public_file = signature_public_file
+
   def upload(self, file_descriptor, directory_key=None, **kw):
     ''' Upload the file to the server.
     If directory_key is None it must only upload to SHACACHE.
@@ -106,9 +114,7 @@ class NetworkcacheClient(object):
       if sha512 is None:
         kw['sha512'] = sha512sum
 
-      signature = kw.pop('signature', None)
-      if signature is None:
-        signature = ''
+      signature = self._getSignatureString()
       data = [kw, signature]
 
       shadir_connection = httplib.HTTPConnection(self.shadir_host,
@@ -158,14 +164,52 @@ class NetworkcacheClient(object):
       raise DirectoryNotFound(result.read())
 
     data_list = json.loads(data)
-    if len(data_list) > 1:
+    if len(data_list) > 1 and self.signature_public_file in _MARKER:
       raise DirectoryNotFound('Too many entries for a given directory. ' \
                'Directory: %s. Entries: %s.' % (directory_key, str(data_list)))
 
-    information_dict, signature = data_list[0]
-    sha512 = information_dict.get('sha512')
+    sha512 = None
+    if self.signature_private_file not in _MARKER:
+      for information_dict, signature in data_list:
+        if self._verifySignature(signature):
+          sha512 = information_dict.get('sha512')
+          break
+      if sha512 is None:
+        raise DirectoryNotFound('Could not find a trustable entry.')
+
+    else:
+      information_dict, signature = data_list[0]
+      sha512 = information_dict.get('sha512')
+
     return self.download(sha512)
 
+  def _getSignatureString(self):
+    """
+      Return the signature based on certification file.
+    """
+    if self.signature_private_file in _MARKER:
+      return ''
+
+    SignEVP = M2Crypto.EVP.load_key(self.signature_private_file)
+    SignEVP.sign_init()
+    SignEVP.sign_update('')
+    StringSignature = SignEVP.sign_final()
+    return StringSignature.encode('base64')
+
+  def _verifySignature(self, signature_string):
+    """
+      Check if the signature is valid.
+    """
+    if self.signature_public_file in _MARKER:
+      return 0
+
+    PubKey = M2Crypto.X509.load_cert(self.signature_public_file)
+    VerifyEVP = M2Crypto.EVP.PKey()
+    VerifyEVP.assign_rsa(PubKey.get_pubkey().get_rsa())
+    VerifyEVP.verify_init()
+    VerifyEVP.verify_update('')
+    return VerifyEVP.verify_final(signature_string.decode('base64'))
+
 
 class DirectoryNotFound(Exception):
   pass