updater: Implement prepare system

As updater is used in environment, which requires it to have certificates
available as fast as possible, add a prepare step and allow to launch it with
--prepare-only switch.

Thanks to this it is possible to run it with configuration file to provide
fallback or master certificates for all slaves without connecting to the
network, thus resulting in fast preparation.

/reviewed-on !3

kedifa/cli.py

@@ -189,12 +189,19 @@ def updater(*args):
     help='Run only once.',
   )
 
+  parser.add_argument(
+    '--prepare-only',
+    action='store_true',
+    help='Only prepares, without using netowrk. Enforces --once, disables '
+         '--on-update, does not use nor lock state file, as it is not used.',
+  )
+
   parsed = parser.parse_args(args)
 
   u = Updater(
     parsed.sleep, parsed.mapping.name, parsed.state, parsed.master_certificate,
     parsed.on_update, parsed.identity.name, parsed.server_ca_certificate.name,
-    parsed.once
+    parsed.once, parsed.prepare_only
   )
   parsed.mapping.close()
   parsed.identity.close()

kedifa/test.py

@@ -1111,42 +1111,42 @@ class KedifaUpdaterMappingTest(KedifaUpdaterMixin, unittest.TestCase):
   def test_updateMapping_empty(self):
     self.setupMapping()
     u = updater.Updater(1, self.mapping, self.state, None, None, None, None,
-                        True)
+                        True, False)
     u.updateMapping()
     self.assertEqual(u.mapping, {})
 
   def test_updateMapping_normal(self):
     self.setupMapping('url file')
     u = updater.Updater(1, self.mapping, self.state, None, None, None, None,
-                        True)
+                        True, False)
     u.updateMapping()
     self.assertEqual(u.mapping, {'file': ('url', None)})
 
   def test_updateMapping_morewhite(self):
     self.setupMapping('url \t file')
     u = updater.Updater(1, self.mapping, self.state, None, None, None, None,
-                        True)
+                        True, False)
     u.updateMapping()
     self.assertEqual(u.mapping, {'file': ('url', None)})
 
   def test_updateMapping_one_empty(self):
     self.setupMapping('url file\n     \n')
     u = updater.Updater(1, self.mapping, self.state, None, None, None, None,
-                        True)
+                        True, False)
     u.updateMapping()
     self.assertEqual(u.mapping, {'file': ('url', None)})
 
   def test_updateMapping_one_not_enough(self):
     self.setupMapping('url file\nbuzz\n')
     u = updater.Updater(1, self.mapping, self.state, None, None, None, None,
-                        True)
+                        True, False)
     u.updateMapping()
     self.assertEqual(u.mapping, {'file': ('url', None)})
 
   def test_updateMapping_with_fallback(self):
     self.setupMapping('url file\nbuzz oink fallback\n')
     u = updater.Updater(1, self.mapping, self.state, None, None, None, None,
-                        True)
+                        True, False)
     u.updateMapping()
     self.assertEqual(
       u.mapping, {'file': ('url', None), 'oink': ('buzz', 'fallback')})
@@ -1154,7 +1154,7 @@ class KedifaUpdaterMappingTest(KedifaUpdaterMixin, unittest.TestCase):
   def test_updateMapping_one_comment(self):
     self.setupMapping('url file\n#buzz uff\n')
     u = updater.Updater(1, self.mapping, self.state, None, None, None, None,
-                        True)
+                        True, False)
     u.updateMapping()
     self.assertEqual(u.mapping, {'file': ('url', None)})
 
@@ -1164,7 +1164,7 @@ class KedifaUpdaterUpdateCertificateTest(
   def setUp(self):
     super(KedifaUpdaterUpdateCertificateTest, self).setUp()
     certificate_file = tempfile.NamedTemporaryFile(
-      dir=self.testdir, delete=False)
+      dir=self.testdir, delete=True)
     certificate_file.close()
     self.certificate_file_name = certificate_file.name
 
@@ -1183,7 +1183,7 @@ class KedifaUpdaterUpdateCertificateTest(
     self.setupMapping(mapping)
     u = updater.Updater(
       1, self.mapping, self.state, '/master/certificate/file', None, None,
-      None, True)
+      None, True, False)
     u.updateMapping()
     u.readState()
     with mock.patch.object(
@@ -1339,29 +1339,111 @@ class KedifaUpdaterUpdateCertificateTest(
     self.assertTrue(update)
     self.assertState({self.certificate_file_name: True})
 
+  def _prepare(self, certificate, master_content, fallback=None):
+    if certificate:
+      with open(self.certificate_file_name, 'w') as fh:
+        fh.write(certificate)
+    fallback_file = None
+    if fallback:
+      fallback_file = tempfile.NamedTemporaryFile(
+        dir=self.testdir, delete=False)
+      fallback_file.write(fallback)
+      fallback_file.close()
+    master_file = '/master/certificate/file'
+    if master_content:
+      master_file = tempfile.NamedTemporaryFile(
+        dir=self.testdir, delete=False)
+      master_file.write(master_content)
+      master_file.close()
+      master_file = master_file.name
+
+    mapping = 'http://example.com %s' % (self.certificate_file_name,)
+    if fallback_file:
+      mapping = '%s %s' % (mapping, fallback_file.name)
+    self.setupMapping(mapping)
+    u = updater.Updater(
+      1, self.mapping, self.state, master_file, None, None,
+      None, True, True)
+    with mock.patch.object(
+      updater.Updater, 'fetchCertificate') as fetchCertificate:
+      with mock.patch.object(
+        updater.Updater, 'writeState') as writeState:
+        u.prepare()
+    writeState.assert_not_called()
+    fetchCertificate.assert_not_called()
+    try:
+      return open(self.certificate_file_name, 'r').read()
+    except IOError:
+      return None
+
+  def test_nocert_nomaster_nofallback(self):
+    certificate = self._prepare(
+      certificate='', master_content=None)
+    self.assertEqual(None, certificate)
+
+  def test_nocert_master_nofallback(self):
+    certificate = self._prepare(
+      certificate='', master_content='master')
+    self.assertEqual('master', certificate)
+
+  def test_nocert_nomaster_fallback(self):
+    certificate = self._prepare(
+      certificate='', master_content=None, fallback='fallback')
+    self.assertEqual('fallback', certificate)
+
+  def test_nocert_master_fallback(self):
+    certificate = self._prepare(
+      certificate='', master_content='master', fallback='fallback')
+    self.assertEqual('fallback', certificate)
+
+  def test_cert_nomaster_nofallback(self):
+    certificate = self._prepare(
+      certificate='cert', master_content=None)
+    self.assertEqual('cert', certificate)
+
+  def test_cert_master_nofallback(self):
+    certificate = self._prepare(
+      certificate='cert', master_content='master')
+    self.assertEqual('cert', certificate)
+
+  def test_cert_nomaster_fallback(self):
+    certificate = self._prepare(
+      certificate='cert', master_content=None, fallback='fallback')
+    self.assertEqual('cert', certificate)
+
+  def test_cert_master_fallback(self):
+    certificate = self._prepare(
+      certificate='cert', master_content='master', fallback='fallback')
+    self.assertEqual('cert', certificate)
+
 
 class KedifaUpdaterLoopTest(
   KedifaUpdaterMixin, unittest.TestCase):
 
   def test(self):
     u = updater.Updater(
-      1, None, self.state, None, None, None, None, True)
+      1, None, self.state, None, None, None, None, True, False)
     lock_file = u.state_lock_file
     os.unlink(self.state)
     self.assertFalse(os.path.exists(lock_file))
     self.assertFalse(os.path.exists(self.state))
     with mock.patch.object(
-      updater.Updater, 'action', return_value=None) as mock_object:
+      updater.Updater, 'prepare') as mock_prepare:
+      with mock.patch.object(
+        updater.Updater, 'action', return_value=None) as mock_action:
         u.loop()
-    mock_object.assert_called()
+    mock_prepare.assert_called()
+    mock_action.assert_called()
     self.assertFalse(os.path.exists(lock_file))
     self.assertFalse(os.path.exists(self.state))
 
   def test_raises(self):
     u = updater.Updater(
-      1, None, self.state, None, None, None, None, True)
+      1, None, self.state, None, None, None, None, True, False)
     lock_file = u.state_lock_file
     self.assertFalse(os.path.exists(lock_file))
+    with mock.patch.object(
+      updater.Updater, 'prepare'):
       with mock.patch.object(
         updater.Updater, 'action', side_effect=Exception()) as mock_object:
         self.assertRaises(Exception, u.loop)
@@ -1370,7 +1452,7 @@ class KedifaUpdaterLoopTest(
 
   def test_lock(self):
     u = updater.Updater(
-      1, None, self.state, None, None, None, None, True)
+      1, None, self.state, None, None, None, None, True, False)
     lock_file = u.state_lock_file
     lock = zc.lockfile.LockFile(lock_file)
     try:
@@ -1385,17 +1467,20 @@ class KedifaUpdaterLoopTest(
 
   def test_infinite(self):
     u = updater.Updater(
-      1, None, self.state, None, None, None, None, False)
+      1, None, self.state, None, None, None, None, False, False)
     lock_file = u.state_lock_file
     os.unlink(self.state)
     self.assertFalse(os.path.exists(lock_file))
     self.assertFalse(os.path.exists(self.state))
     with mock.patch.object(
-      updater.Updater, 'action', return_value=None) as mock_object:
+      updater.Updater, 'prepare', return_value=None) as mock_prepare:
+      with mock.patch.object(
+        updater.Updater, 'action', return_value=None) as mock_action:
         with mock.patch.object(
           updater.time, 'sleep', side_effect=ValueError('timer')) as timer:
           self.assertRaises(ValueError, u.loop)
     timer.assert_called_with(1)
-    mock_object.assert_called()
+    mock_prepare.assert_called()
+    mock_action.assert_called()
     self.assertFalse(os.path.exists(lock_file))
     self.assertFalse(os.path.exists(self.state))

kedifa/updater.py

@@ -9,7 +9,8 @@ import zc.lockfile
 
 class Updater(object):
   def __init__(self, sleep, mapping_file, state_file, master_certificate_file,
-               on_update, identity_file, server_ca_certificate_file, once):
+               on_update, identity_file, server_ca_certificate_file, once,
+               prepare_only):
     self.sleep = sleep
     self.mapping_file = mapping_file
     self.state_file = state_file
@@ -19,6 +20,7 @@ class Updater(object):
     self.identity_file = identity_file
     self.server_ca_certificate_file = server_ca_certificate_file
     self.once = once
+    self.prepare_only = prepare_only
 
   def updateMapping(self):
     self.mapping = {}
@@ -38,6 +40,10 @@ class Updater(object):
         else:
           print 'Line %r is incorrect' % (line,)
           continue
+        if certificate in self.mapping:
+          print 'Line %r is incorrect, duplicated certificate %r' % (
+            line, certificate)
+          raise ValueError
         self.mapping[certificate] = (url, fallback)
 
   def fetchCertificate(self, url, certificate_file):
@@ -116,9 +122,41 @@ class Updater(object):
     with open(self.state_file, 'w') as fh:
       json.dump(self.state_dict, fh, indent=2)
 
+  def prepare(self):
+    self.updateMapping()
+    prepare_mapping = self.mapping.copy()
+    _, master_certificate_file_fallback = prepare_mapping.pop(
+      self.master_certificate_file, (None, None))
+
+    # update master certificate from fallback
+    if self.master_certificate_file:
+      if not os.path.exists(self.master_certificate_file):
+        if master_certificate_file_fallback and os.path.exists(
+          master_certificate_file_fallback):
+          open(self.master_certificate_file, 'w').write(
+            open(master_certificate_file_fallback, 'r').read()
+          )
+          print 'Prepare: Used %r for %r' % (
+            master_certificate_file_fallback, self.master_certificate_file)
+
+    master_content = None
+    if self.master_certificate_file and os.path.exists(
+      self.master_certificate_file):
+      master_content = open(self.master_certificate_file, 'r').read()
+
+    for certificate, (_, fallback) in prepare_mapping.items():
+      if os.path.exists(certificate):
+        continue
+      if fallback and os.path.exists(fallback):
+        open(certificate, 'w').write(open(fallback, 'r').read())
+        print 'Prepare: Used %r for %r' % (fallback, certificate)
+      elif master_content:
+        open(certificate, 'w').write(master_content)
+        print 'Prepare: Used %r for %r' % (
+          self.master_certificate_file, certificate)
+
   def action(self):
     self.readState()
-    self.updateMapping()
     updated = False
 
     if self.master_certificate_file in self.mapping:
@@ -147,24 +185,28 @@ class Updater(object):
   def loop(self):
     while True:
       try:
+        if not self.prepare_only:
           lock = zc.lockfile.LockFile(self.state_lock_file)
       except zc.lockfile.LockError as e:
         print e,
-        if self.once:
+        if self.once or self.prepare_only:
           print '...exiting.'
           sys.exit(1)
         else:
           print "...will try again later."
       else:
         try:
+          self.prepare()
+          if not self.prepare_only:
             self.action()
         finally:
+          if not self.prepare_only:
             lock.close()
             try:
               os.unlink(self.state_lock_file)
             except Exception as e:
               print 'Problem while unlinking %r' % (self.state_lock_file,)
-      if self.once:
+      if self.once or self.prepare_only:
         break
       print 'Sleeping for %is' % (self.sleep,)
       time.sleep(self.sleep)