source file: /home/buildslave/tahoe/edgy/build/src/allmydata/key_generator.py
file stats: 86 lines, 82 executed: 95.3% covered
   1. 
   2. import os
   3. import time
   4. 
   5. import foolscap
   6. from zope.interface import implements
   7. from twisted.internet import reactor
   8. from twisted.application import service
   9. from twisted.python import log
  10. 
  11. from pycryptopp.publickey import rsa
  12. from allmydata.interfaces import RIKeyGenerator
  13. 
  14. class KeyGenerator(service.MultiService, foolscap.Referenceable):
  15.     implements(RIKeyGenerator)
  16. 
  17.     pool_size = 16 # no. keys to keep on hand in the pool
  18.     pool_refresh_delay = 6 # no. sec to wait after a fetch before generating new keys
  19.     verbose = False
  20. 
  21.     def __init__(self, default_key_size=2048):
  22.         service.MultiService.__init__(self)
  23.         self.keypool = []
  24.         self.last_fetch = 0
  25.         self.default_key_size = default_key_size
  26. 
  27.     def startService(self):
  28.         self.timer = reactor.callLater(0, self.maybe_refill_pool)
  29.         return service.MultiService.startService(self)
  30. 
  31.     def stopService(self):
  32.         if self.timer.active():
  33.             self.timer.cancel()
  34.         return service.MultiService.stopService(self)
  35. 
  36.     def __repr__(self):
  37.         return '<KeyGenerator[%s]>' % (len(self.keypool),)
  38. 
  39.     def vlog(self, msg):
  40.         if self.verbose:
  41.             log.msg(msg)
  42. 
  43.     def reset_timer(self):
  44.         self.last_fetch = time.time()
  45.         if self.timer.active():
  46.             self.timer.reset(self.pool_refresh_delay)
  47.         else:
  48.             self.timer = reactor.callLater(self.pool_refresh_delay, self.maybe_refill_pool)
  49. 
  50.     def maybe_refill_pool(self):
  51.         now = time.time()
  52.         if self.last_fetch + self.pool_refresh_delay < now:
  53.             self.vlog('%s refilling pool' % (self,))
  54.             while len(self.keypool) < self.pool_size:
  55.                 self.keypool.append(self.gen_key(self.default_key_size))
  56.         else:
  57.             self.vlog('%s not refilling pool' % (self,))
  58.             reactor.callLater(1, self.maybe_refill_pool)
  59. 
  60.     def gen_key(self, key_size):
  61.         self.vlog('%s generating key size %s' % (self, key_size, ))
  62.         signer = rsa.generate(key_size)
  63.         verifier = signer.get_verifying_key()
  64.         return verifier.serialize(), signer.serialize()
  65. 
  66.     def remote_get_rsa_key_pair(self, key_size):
  67.         self.vlog('%s remote_get_key' % (self,))
  68.         if key_size != self.default_key_size or not self.keypool:
  69.             key = self.gen_key(key_size)
  70.             self.reset_timer()
  71.             return key
  72.         else:
  73.             self.reset_timer()
  74.             return self.keypool.pop()
  75. 
  76. class KeyGeneratorService(service.MultiService):
  77.     furl_file = 'key_generator.furl'
  78. 
  79.     def __init__(self, basedir='.', display_furl=True, default_key_size=2048):
  80.         service.MultiService.__init__(self)
  81.         self.basedir = basedir
  82.         self.tub = foolscap.Tub(certFile=os.path.join(self.basedir, 'key_generator.pem'))
  83.         self.tub.setServiceParent(self)
  84.         self.key_generator = KeyGenerator(default_key_size=default_key_size)
  85.         self.key_generator.setServiceParent(self)
  86. 
  87.         portnum = self.get_portnum()
  88.         self.listener = self.tub.listenOn(portnum or 'tcp:0')
  89.         d = self.tub.setLocationAutomatically()
  90.         if portnum is None:
  91.             d.addCallback(self.save_portnum)
  92.         d.addCallback(self.tub_ready, display_furl)
  93.         d.addErrback(log.err)
  94. 
  95.     def get_portnum(self):
  96.         portnumfile = os.path.join(self.basedir, 'portnum')
  97.         if os.path.exists(portnumfile):
  98.             return file(portnumfile, 'rb').read().strip()
  99. 
 100.     def save_portnum(self, junk):
 101.         portnum = self.listener.getPortnum()
 102.         portnumfile = os.path.join(self.basedir, 'portnum')
 103.         file(portnumfile, 'wb').write('%d\n' % (portnum,))
 104. 
 105.     def tub_ready(self, junk, display_furl):
 106.         kgf = os.path.join(self.basedir, self.furl_file)
 107.         self.keygen_furl = self.tub.registerReference(self.key_generator, furlFile=kgf)
 108.         if display_furl:
 109.             print 'key generator at:', self.keygen_furl