swarming: Add server-side implementation for supplemental bot_config

appengine/swarming/server/bot_groups_config.py

 # Copyright 2016 The LUCI Authors. All rights reserved.
 # Use of this source code is governed under the Apache License, Version 2.0
 # that can be found in the LICENSE file.
 
 """Functions to fetch and interpret bots.cfg file with list of bot groups."""
 
 import collections
 import hashlib
 import logging
+import os
 import re
 
 from components import auth
 from components import config
 from components import utils
 from components.config import validation
 
 from proto import bots_pb2
 from server import config as local_config
 
@@ skipping 20 matching lines @@
   'ip_whitelist',
 
   # Tuple with emails of bot owners.
   'owners',
 
   # Dict {key => list of values}. Always contains all the keys specified by
   # 'trusted_dimensions' set in BotsCfg. If BotGroup doesn't define some
   # dimension from that set, the list of value for it will be empty. Key and
   # values are unicode strings.
   'dimensions',
+
+  # Name of the supplemental bot_config.py to inject to the bot during
+  # handshake.
+  'bot_config_script',
+
+  # Content of the supplemental bot_config.py to inject to the bot during
+  # handshake.
+  'bot_config_script_content',
 ])
 
 
 # Post-processed and validated read-only form of bots.cfg config. Its structure
 # is optimized for fast lookup of BotGroupConfig by bot_id.
 _BotGroups = collections.namedtuple('_BotGroups', [
   'direct_matches', # dict bot_id => BotGroupConfig
   'prefix_matches', # list of pairs (bot_id_prefix, BotGroupConfig)
   'machine_types',  # dict machine_type.name => BotGroupConfig
   'default_group',  # fallback BotGroupConfig or None if not defined
 ])
 
 
 # Default config to use on unconfigured server.
 def _default_bot_groups():
   return _BotGroups(
     direct_matches={},
     prefix_matches=[],
     machine_types={},
     default_group=BotGroupConfig(
         version='default',
         require_luci_machine_token=False,
         require_service_account=None,
         ip_whitelist=auth.bots_ip_whitelist(),
         owners=(),
-        dimensions={}))
+        dimensions={},
+        bot_config_script='',
+        bot_config_script_content=''))
 
 
 def _gen_version(fields):
   """Looks at BotGroupConfig fields and derives a digest that summarizes them.
 
   This digest is going to be sent to the bot in /handshake, and bot would
   include it in its state (and thus send it with each /poll). If server detects
   that the bot is using older version of the config, it would ask the bot
   to restart.
 
@@ skipping 48 matching lines @@
     # In validated config 'dim_kv_pair' is always 'key:value', but be cautious.
     parts = unicode(dim_kv_pair).split(':', 1)
     if len(parts) != 2:
       logging.error('Invalid dimension in bots.cfg - "%s"', dim_kv_pair)
       continue
     k, v = parts[0], parts[1]
     dimensions.setdefault(k, set()).add(v)
 
   auth_cfg = msg.auth or bots_pb2.BotAuth()
 
+  content = ''
+  if msg.bot_config_script:
+    rev, content = config.get_self_config(
+        'scripts/' + msg.bot_config_script,
+        store_last_good=True)
+    if not rev or not content:
+      # The entry is invalid. It points to a non existing file. It could be
+      # because of a typo in the file name. An empty file is an invalid file,
+      # log an error to alert the admins.
+      logging.error(
+          'Configuration referenced non existing bot_config file %r\n%s',
+          msg.bot_config_script, msg)
   return _make_bot_group_config(
     require_luci_machine_token=auth_cfg.require_luci_machine_token,
     require_service_account=auth_cfg.require_service_account,
     ip_whitelist=auth_cfg.ip_whitelist,
     owners=tuple(msg.owners),
-    dimensions={k: sorted(v) for k, v in dimensions.iteritems()})
+    dimensions={k: sorted(v) for k, v in dimensions.iteritems()},
+    bot_config_script=msg.bot_config_script or '',
+    bot_config_script_content=content or '')
 
 
 def _expand_bot_id_expr(expr):
   """Expands string with bash-like sets (if they are there).
 
   E.g. takes "vm{1..3}-m1" and yields "vm1-m1", "vm2-m1", "vm3-m1". Also
   supports list syntax ({1,2,3}). Either one should be used, but not both, e.g.
   following WILL NOT work: {1..3,4,5}.
 
   Yields original string if it doesn't have '{...}' section.
@@ skipping 226 matching lines @@
         try:
           auth.Identity(auth.IDENTITY_USER, own)
         except ValueError:
           ctx.error('invalid owner email "%s"', own)
 
       # Validate 'dimensions'.
       for dim in entry.dimensions:
         if not local_config.validate_flat_dimension(dim):
           ctx.error('bad dimension %r', dim)
 
+      # Validate 'bot_config_script': the supplemental bot_config.py.
+      if entry.bot_config_script:
+        # Another check in bot_code.py confirms that the script itself is valid
+        # python.
+        if not entry.bot_config_script.endswith('.py'):
+          ctx.error('Invalid bot_config_script name: must end with .py')
+        if os.path.basename(entry.bot_config_script) != entry.bot_config_script:
+          ctx.error(
+              'Invalid bot_config_script name: must not contain path entry')
+        # We can't validate that the file exists here. It'll fail in
+        # _bot_group_proto_to_tuple() which is called by _fetch_bot_groups() and
+        # cached for 60 seconds.
+
   # Now verify bot_id_prefix is never a prefix of other prefix. It causes
   # ambiguities.
   for smaller, s_idx in bot_id_prefixes.iteritems():
     for larger, l_idx in bot_id_prefixes.iteritems():
       if smaller == larger:
         continue # we've already checked prefixes have no duplicated
       if larger.startswith(smaller):
         ctx.error(
             'bot_id_prefix "%s", defined in group #%d, is subprefix of "%s", '
             'defined in group #%d; it makes group assigned for bots with '
             'prefix "%s" ambigious', smaller, s_idx, larger, l_idx, larger)