1# Copyright 2015 The Chromium OS Authors. All rights reserved.
2# Use of this source code is governed by a BSD-style license that can be
3# found in the LICENSE file.
4#
5# This implements the APIs defined at
6# https://github.com/luci/luci-py/blob/master/appengine/
7# swarming/swarming_bot/bot_config.py
8
9
10"""This file is meant to be overriden by the server's specific copy.
11
12You can upload a new version via /restricted/upload/bot_config.
13
14There's 3 types of functions in this file:
15  - get_*() to return properties to describe this bot.
16  - on_*() as hooks based on events happening on the bot.
17  - setup_*() to setup global state on the host.
18
19This file shouldn't import from other scripts in this directory except
20os_utilities which is guaranteed to be usable as an API. It's fine to import
21from stdlib.
22
23Set the environment variable SWARMING_LOAD_TEST=1 to disable the use of
24server-provided bot_config.py. This permits safe load testing.
25
26TODO(fdeng):
27    Restrict the command to run_suite/abort_suite
28"""
29
30import json
31import re
32import os
33
34from api import os_utilities
35
36# Unused argument 'bot' - pylint: disable=W0613
37
38
39CMD_WHITELIST = {'/usr/local/autotest/site_utils/run_suite.py',
40                 '/usr/local/autotest/site_utils/abort_suite.py'}
41
42
43def get_dimensions(bot=None):
44    """Returns dict with the bot's dimensions.
45
46    The dimensions are what are used to select the bot that can run each task.
47
48    By default, the bot id will be automatically selected based on
49    the hostname with os_utilities.get_dimensions(). This method
50    overrides the default id returned by os_utilities.get_dimensions().
51
52    Assume the bot's working directory is like BOT_ROOT/bot_23/
53    we will parse the id "23" from the directory name and append it to the
54    hostname to form the bot id. so the bot id would look like
55    chromeos-server31-23
56
57    See https://github.com/luci/luci-py/blob/master/appengine/
58    swarming/doc/Magic-Values.md
59
60    @returns: Dict with the bot's dimentions.
61
62    """
63    d = os_utilities.get_dimensions()
64    m = re.match('.*/bot_([\d]+).*', os.getcwd())
65    suffix = ''
66    if m:
67        suffix = '-'+ m.group(1)
68    d[u'id'] = [os_utilities.get_hostname_short() + suffix]
69    return d
70
71
72def get_state(bot=None):
73    """Returns dict with a state of the bot reported to the server with each poll.
74
75    It is only for dynamic state that changes while bot is running for information
76    for the sysadmins.
77
78    The server can not use this state for immediate scheduling purposes (use
79    'dimensions' for that), but it can use it for maintenance and bookkeeping
80    tasks.
81
82    See https://github.com/luci/luci-py/blob/master/appengine/
83    swarming/doc/Magic-Values.md
84
85    """
86    return os_utilities.get_state()
87
88
89### Hooks
90
91
92def on_before_task(bot):
93    """Hook function called before running a task.
94
95    It shouldn't do much, since it can't cancel the task so it shouldn't do
96    anything too fancy.
97
98    @param bot: bot.Bot instance.
99    """
100    # TODO(fdeng): it is possible that the format gets updated
101    # without warning. It would be better to find a long term solution.
102    work_dir = os.path.join(bot.base_dir, 'work')
103    path = os.path.join(work_dir, 'task_runner_in.json')
104    manifest = {}
105    with open(path) as f:
106        manifest = json.load(f)
107    full_command = manifest.get('command')
108    if full_command and not full_command[0] in CMD_WHITELIST:
109        # override the command with a safe "echo"
110        manifest['command'] = ['echo', '"Command not allowed"']
111        manifest['data'] = []
112        with open(path, 'wb') as f:
113            f.write(json.dumps(manifest))
114        raise Exception('Command not allowed: %s' % full_command)
115
116
117### Setup
118
119
120def setup_bot(bot):
121    """Does one time initialization for this bot.
122
123    Returns True if it's fine to start the bot right away. Otherwise, the calling
124    script should exit.
125
126    Example: making this script starts automatically on user login via
127    os_utilities.set_auto_startup_win() or os_utilities.set_auto_startup_osx().
128
129    @param bot: bot.Bot instance.
130
131    @returns: Boolean. See above.
132
133    """
134    return True
135