1#!/usr/bin/env python
2#
3# Copyright (C) 2017 The Android Open Source Project
4#
5# Licensed under the Apache License, Version 2.0 (the "License");
6# you may not use this file except in compliance with the License.
7# You may obtain a copy of the License at
8#
9#      http://www.apache.org/licenses/LICENSE-2.0
10#
11# Unless required by applicable law or agreed to in writing, software
12# distributed under the License is distributed on an "AS IS" BASIS,
13# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14# See the License for the specific language governing permissions and
15# limitations under the License.
16#
17
18import argparse
19import json
20import logging
21import socket
22import time
23import threading
24import sys
25
26from host_controller import console
27from host_controller import tfc_host_controller
28from host_controller.build import build_provider_pab
29from host_controller.tfc import tfc_client
30from host_controller.vti_interface import vti_endpoint_client
31from host_controller.tradefed import remote_client
32from vts.utils.python.os import env_utils
33
34_ANDROID_BUILD_TOP = "ANDROID_BUILD_TOP"
35_SECONDS_PER_UNIT = {
36    "m": 60,
37    "h": 60 * 60,
38    "d": 60 * 60 * 24
39}
40
41
42def _ParseInterval(interval_str):
43    """Parses string to time interval.
44
45    Args:
46        interval_str: string, a floating-point number followed by time unit.
47
48    Returns:
49        float, the interval in seconds.
50
51    Raises:
52        ValueError if the argument format is wrong.
53    """
54    if not interval_str:
55        raise ValueError("Argument is empty.")
56
57    unit = interval_str[-1]
58    if unit not in _SECONDS_PER_UNIT:
59        raise ValueError("Unknown unit: %s" % unit)
60
61    interval = float(interval_str[:-1])
62    if interval < 0:
63        raise ValueError("Invalid time interval: %s" % interval)
64
65    return interval * _SECONDS_PER_UNIT[unit]
66
67
68def _ScriptLoop(hc_console, script_path, loop_interval):
69    """Runs a console script repeatedly.
70
71    Args:
72        hc_console: the host controller console.
73        script_path: string, the path to the script.
74        loop_interval: float or integer, the interval in seconds.
75    """
76    next_start_time = time.time()
77    while hc_console.ProcessScript(script_path):
78        if loop_interval == 0:
79            continue
80        current_time = time.time()
81        skip_cnt = (current_time - next_start_time) // loop_interval
82        if skip_cnt >= 1:
83            logging.warning("Script execution time is longer than loop "
84                            "interval. Skip %d iteration(s).", skip_cnt)
85        next_start_time += (skip_cnt + 1) * loop_interval
86        if next_start_time - current_time >= 0:
87            time.sleep(next_start_time - current_time)
88        else:
89            logging.error("Unexpected timestamps: current=%f, next=%f",
90                          current_time, next_start_time)
91
92
93def main():
94    """Parses arguments and starts console."""
95    parser = argparse.ArgumentParser()
96    parser.add_argument("--config-file",
97                        default=None,
98                        type=argparse.FileType('r'),
99                        help="The configuration file in JSON format")
100    parser.add_argument("--poll", action="store_true",
101                        help="Disable console and start host controller "
102                             "threads polling TFC.")
103    parser.add_argument("--use-tfc", action="store_true",
104                        help="Enable TFC (TradeFed Cluster).")
105    parser.add_argument("--vti",
106                        default=None,
107                        help="The base address of VTI endpoint APIs")
108    parser.add_argument("--script",
109                        default=None,
110                        help="The path to a script file in .py format")
111    parser.add_argument("--serial",
112                        default=None,
113                        help="The default serial numbers for flashing and "
114                             "testing in the console. Multiple serial numbers "
115                             "are separated by comma.")
116    parser.add_argument("--loop",
117                        default=None,
118                        metavar="INTERVAL",
119                        type=_ParseInterval,
120                        help="The interval of repeating the script. "
121                             "The format is a float followed by unit which is "
122                             "one of 'm' (minute), 'h' (hour), and 'd' (day). "
123                             "If this option is unspecified, the script will "
124                             "be processed once.")
125    parser.add_argument("--console", action="store_true",
126                        help="Whether to start a console after processing "
127                             "a script.")
128    parser.add_argument("--password",
129                        default=None,
130                        help="Password string to pass to the prompt "
131                             "when running certain command as root previlege.")
132    parser.add_argument("--flash",
133                        default=None,
134                        help="GCS URL to an img package. Fetches and flashes "
135                             "the device(s) given as the '--serial' flag.")
136    args = parser.parse_args()
137    if args.config_file:
138        config_json = json.load(args.config_file)
139    else:
140        config_json = {}
141        config_json["log_level"] = "DEBUG"
142        config_json["hosts"] = []
143        host_config = {}
144        host_config["cluster_ids"] = ["local-cluster-1",
145                                      "local-cluster-2"]
146        host_config["lease_interval_sec"] = 30
147        config_json["hosts"].append(host_config)
148
149    env_vars = env_utils.SaveAndClearEnvVars([_ANDROID_BUILD_TOP])
150
151    root_logger = logging.getLogger()
152    root_logger.setLevel(getattr(logging, config_json["log_level"]))
153
154    if args.vti:
155        vti_endpoint = vti_endpoint_client.VtiEndpointClient(args.vti)
156    else:
157        vti_endpoint = None
158
159    tfc = None
160    if args.use_tfc:
161        if args.config_file:
162            tfc = tfc_client.CreateTfcClient(
163                    config_json["tfc_api_root"],
164                    config_json["service_key_json_path"],
165                    api_name=config_json["tfc_api_name"],
166                    api_version=config_json["tfc_api_version"],
167                    scopes=config_json["tfc_scopes"])
168        else:
169            logging.warning("WARN: If --use_tfc is set, --config_file argument "
170                            "value must be provided. Starting without TFC.")
171
172    pab = build_provider_pab.BuildProviderPAB()
173
174    hosts = []
175    for host_config in config_json["hosts"]:
176        cluster_ids = host_config["cluster_ids"]
177        # If host name is not specified, use local host.
178        hostname = host_config.get("hostname", socket.gethostname())
179        port = host_config.get("port", remote_client.DEFAULT_PORT)
180        cluster_ids = host_config["cluster_ids"]
181        remote = remote_client.RemoteClient(hostname, port)
182        host = tfc_host_controller.HostController(remote, tfc, hostname,
183                                                  cluster_ids)
184        hosts.append(host)
185        if args.poll:
186            lease_interval_sec = host_config["lease_interval_sec"]
187            host_thread = threading.Thread(target=host.Run,
188                                           args=(lease_interval_sec,))
189            host_thread.daemon = True
190            host_thread.start()
191
192    if args.poll:
193        while True:
194            sys.stdin.readline()
195    else:
196        main_console = console.Console(vti_endpoint, tfc, pab, hosts,
197                                       vti_address=args.vti,
198                                       password=args.password)
199        if args.vti:
200            main_console.StartJobThreadAndProcessPool()
201        else:
202            logging.warning("vti address is not set. example : "
203                            "$ run --vti=<url>")
204
205        try:
206            if args.serial:
207                main_console.SetSerials(args.serial.split(","))
208            if args.script:
209                if args.loop is None:
210                    main_console.ProcessScript(args.script)
211                else:
212                    _ScriptLoop(main_console, args.script, args.loop)
213
214                if args.console:
215                    main_console.cmdloop()
216            elif args.flash:
217                main_console.FlashImgPackage(args.flash)
218            else:  # if not script, the default is console mode.
219                main_console.cmdloop()
220        finally:
221            main_console.TearDown()
222
223    env_utils.RestoreEnvVars(env_vars)
224
225
226if __name__ == "__main__":
227    main()
228