1#
2# Copyright (C) 2018 The Android Open Source Project
3#
4# Licensed under the Apache License, Version 2.0 (the 'License');
5# you may not use this file except in compliance with the License.
6# You may obtain a copy of the License at
7#
8#      http://www.apache.org/licenses/LICENSE-2.0
9#
10# Unless required by applicable law or agreed to in writing, software
11# distributed under the License is distributed on an 'AS IS' BASIS,
12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13# See the License for the specific language governing permissions and
14# limitations under the License.
15#
16
17import httplib2
18import itertools
19import logging
20import os
21import socket
22import threading
23import time
24
25from googleapiclient import errors
26from google.protobuf import text_format
27
28from host_controller import common
29from host_controller.command_processor import base_command_processor
30from host_controller.console_argument_parser import ConsoleArgumentError
31from host_controller.tradefed import remote_operation
32
33from vti.test_serving.proto import TestLabConfigMessage_pb2 as LabCfgMsg
34from vti.test_serving.proto import TestScheduleConfigMessage_pb2 as SchedCfgMsg
35
36
37class CommandConfig(base_command_processor.BaseCommandProcessor):
38    """Command processor for config command.
39
40    Attributes:
41        arg_parser: ConsoleArgumentParser object, argument parser.
42        console: cmd.Cmd console object.
43        command: string, command name which this processor will handle.
44        command_detail: string, detailed explanation for the command.
45        schedule_thread: dict containing threading.Thread instances(s) that
46                         update schedule info regularly.
47    """
48
49    command = "config"
50    command_detail = "Specifies a global config type to monitor."
51
52    def UpdateConfig(self, account_id, branch, targets, config_type, method,
53                     update_build, clear_schedule, clear_labinfo):
54        """Updates the global configuration data.
55
56        Args:
57            account_id: string, Partner Android Build account_id to use.
58            branch: string, branch to grab the artifact from.
59            targets: string, a comma-separate list of build target product(s).
60            config_type: string, config type (`prod` or `test').
61            method: string, HTTP method for fetching.
62            update_build: boolean, indicating whether to upload build info.
63            clear_schedule: bool, True to clear all schedule data exist on the
64                            scheduler
65            clear_labinfo: bool, True to clear all lab data exist on the
66                           scheduler
67        """
68        for target in targets.split(","):
69            fetch_path = self.FetchConfig(
70                account_id=account_id,
71                branch=branch,
72                target=target,
73                config_type=config_type,
74                method=method)
75            if fetch_path:
76                self.UploadConfig(
77                    path=fetch_path,
78                    update_build=update_build,
79                    clear_schedule=clear_schedule,
80                    clear_labinfo=clear_labinfo)
81
82    def FetchConfig(self, account_id, branch, target, config_type, method):
83        """Fetches config files from the PAB build provider.
84
85        Args:
86            account_id: string, Partner Android Build account_id to use.
87            branch: string, branch to grab the artifact from.
88            target: string, build target.
89            config_type: string, config type (`prod` or `test').
90            method: string, HTTP method for fetching.
91
92        Returns:
93            string, a path to the temp directory where config files are stored.
94        """
95        path = ""
96        self.console._build_provider["pab"].Authenticate()
97        try:
98            listed_builds = self.console._build_provider["pab"].GetBuildList(
99                account_id=account_id,
100                branch=branch,
101                target=target,
102                page_token="",
103                max_results=1,
104                method="GET")
105        except ValueError as e:
106            logging.exception(e)
107            return path
108
109        if listed_builds and len(listed_builds) > 0:
110            listed_build = listed_builds[0]
111            if listed_build["successful"]:
112                device_images, test_suites, artifacts, configs = (
113                    self.console._build_provider["pab"].GetArtifact(
114                        account_id=account_id,
115                        branch=branch,
116                        target=target,
117                        artifact_name=(
118                            "vti-global-config-%s.zip" % config_type),
119                        build_id=listed_build["build_id"],
120                        method=method))
121                path = os.path.dirname(configs[config_type])
122
123        return path
124
125    def UploadConfig(self, path, update_build, clear_schedule, clear_labinfo):
126        """Uploads configs to VTI server.
127
128        Args:
129            path: string, a path where config files are stored.
130            update_build: boolean, indicating whether to upload build info.
131            clear_schedule: bool, True to clear all schedule data exist on the
132                            scheduler
133            clear_labinfo: bool, True to clear all lab data exist on the
134                           scheduler
135        """
136        schedules_pbs = []
137        lab_pbs = []
138        for root, dirs, files in os.walk(path):
139            for config_file in files:
140                full_path = os.path.join(root, config_file)
141                try:
142                    if config_file.endswith(".schedule_config"):
143                        with open(full_path, "r") as fd:
144                            context = fd.read()
145                            sched_cfg_msg = SchedCfgMsg.ScheduleConfigMessage()
146                            text_format.Merge(context, sched_cfg_msg)
147                            schedules_pbs.append(sched_cfg_msg)
148                            logging.info(sched_cfg_msg.manifest_branch)
149                    elif config_file.endswith(".lab_config"):
150                        with open(full_path, "r") as fd:
151                            context = fd.read()
152                            lab_cfg_msg = LabCfgMsg.LabConfigMessage()
153                            text_format.Merge(context, lab_cfg_msg)
154                            lab_pbs.append(lab_cfg_msg)
155                except text_format.ParseError as e:
156                    logging.error("ERROR: Config parsing error %s", e)
157        if update_build:
158            commands = self.GetBuildCommands(schedules_pbs)
159            if commands:
160                for command in commands:
161                    ret = self.console.onecmd(command)
162                    if ret == False:
163                        break
164        self.console._vti_endpoint_client.UploadScheduleInfo(
165            schedules_pbs, clear_schedule)
166        self.console._vti_endpoint_client.UploadLabInfo(lab_pbs, clear_labinfo)
167
168    def UpdateConfigLoop(self, account_id, branch, target, config_type, method,
169                         update_build, update_interval, clear_schedule,
170                         clear_labinfo):
171        """Regularly updates the global configuration.
172
173        Args:
174            account_id: string, Partner Android Build account_id to use.
175            branch: string, branch to grab the artifact from.
176            targets: string, a comma-separate list of build target product(s).
177            config_type: string, config type (`prod` or `test').
178            method: string, HTTP method for fetching.
179            update_build: boolean, indicating whether to upload build info.
180            update_interval: int, number of seconds before repeating
181            clear_schedule: bool, True to clear all schedule data exist on the
182                            scheduler
183            clear_labinfo: bool, True to clear all lab data exist on the
184                           scheduler
185        """
186        thread = threading.currentThread()
187        while getattr(thread, 'keep_running', True):
188            try:
189                self.UpdateConfig(account_id, branch, target, config_type,
190                                  method, update_build, clear_schedule,
191                                  clear_labinfo)
192            except (socket.error, remote_operation.RemoteOperationException,
193                    httplib2.HttpLib2Error, errors.HttpError) as e:
194                logging.exception(e)
195            time.sleep(update_interval)
196
197    def GetBuildCommands(self, schedule_pbs):
198        """Generates a list of build commands with given schedules.
199
200        Args:
201            schedule_pbs: a list of TestScheduleConfig protobuf messages.
202
203        Returns:
204            a list of build command strings
205        """
206        attrs = {}
207        attrs["device"] = [
208            "build_storage_type", "manifest_branch", "pab_account_id",
209            "require_signed_device_build", "name"
210        ]
211        attrs["gsi"] = [
212            "gsi_storage_type", "gsi_branch", "gsi_pab_account_id",
213            "gsi_build_target"
214        ]
215        attrs["test"] = [
216            "test_storage_type", "test_branch", "test_pab_account_id",
217            "test_build_target"
218        ]
219
220        class BuildInfo(object):
221            """A build information class."""
222
223            def __init__(self, _build_type):
224                if _build_type in attrs:
225                    for attribute in attrs[_build_type]:
226                        setattr(self, attribute, "")
227
228            def __eq__(self, compare):
229                return self.__dict__ == compare.__dict__
230
231        build_commands = []
232        if not schedule_pbs:
233            return build_commands
234
235        # parses the given protobuf and stores as BuildInfo object.
236        builds = {"device": [], "gsi": [], "test": []}
237        for pb in schedule_pbs:
238            for build_target in pb.build_target:
239                build_type = "device"
240                device = BuildInfo(build_type)
241                for attr in attrs[build_type]:
242                    if hasattr(pb, attr):
243                        setattr(device, attr, getattr(pb, attr, None))
244                    elif hasattr(build_target, attr):
245                        setattr(device, attr, getattr(build_target, attr,
246                                                      None))
247                if not [x for x in builds[build_type] if x == device]:
248                    builds[build_type].append(device)
249                for test_schedule in build_target.test_schedule:
250                    build_type = "gsi"
251                    gsi = BuildInfo(build_type)
252                    for attr in attrs[build_type]:
253                        if hasattr(test_schedule, attr):
254                            setattr(gsi, attr,
255                                    getattr(test_schedule, attr, None))
256                    if not [x for x in builds[build_type] if x == gsi]:
257                        builds[build_type].append(gsi)
258
259                    build_type = "test"
260                    test = BuildInfo(build_type)
261                    for attr in attrs[build_type]:
262                        if hasattr(test_schedule, attr):
263                            setattr(test, attr,
264                                    getattr(test_schedule, attr, None))
265                    if not [x for x in builds[build_type] if x == test]:
266                        builds[build_type].append(test)
267
268        # groups by artifact, branch, and account id, and builds a command.
269        for artifact in attrs:
270            load_attrs = attrs[artifact]
271            if artifact == "device":
272                storage_type_text = "build_storage_type"
273            else:
274                storage_type_text = "" + artifact + "_storage_type"
275            pab_builds = [
276                x for x in builds[artifact]
277                if getattr(x, storage_type_text) ==
278                SchedCfgMsg.BUILD_STORAGE_TYPE_PAB
279            ]
280            pab_builds.sort(key=lambda x: tuple([getattr(x, attribute)
281                                                 for attribute in load_attrs]))
282            groups = [list(g) for k, g in itertools.groupby(
283                pab_builds, lambda x: tuple([getattr(x, attribute)
284                                             for attribute
285                                             in load_attrs[1:-1]]))]
286            for group in groups:
287                command = ("build --artifact-type={} --method=GET "
288                           "--noauth_local_webserver=True --update=single".
289                           format(artifact))
290                if artifact == "device":
291                    if group[0].manifest_branch:
292                        command += " --branch={}".format(
293                            group[0].manifest_branch)
294                    else:
295                        logging.debug(
296                            "Device manifest branch is a mandatory field.")
297                        continue
298                    if group[0].pab_account_id:
299                        command += " --account_id={}".format(
300                            group[0].pab_account_id)
301                    if group[0].require_signed_device_build:
302                        command += " --verify-signed-build=True"
303                    targets = ",".join([x.name for x in group if x.name])
304                    if targets:
305                        command += " --target={}".format(targets)
306                        build_commands.append(command)
307                else:
308                    if getattr(group[0], "" + artifact + "_branch"):
309                        command += " --branch={}".format(
310                            getattr(group[0], "" + artifact + "_branch"))
311                    else:
312                        logging.debug(
313                            "{} branch is a mandatory field.".format(artifact))
314                        continue
315                    if getattr(group[0], "" + artifact + "_pab_account_id"):
316                        command += " --account_id={}".format(
317                            getattr(group[0],
318                                    "" + artifact + "_pab_account_id"))
319                    targets = ",".join([
320                        getattr(x, "" + artifact + "_build_target")
321                        for x in group
322                        if getattr(x, "" + artifact + "_build_target")
323                    ])
324                    if targets:
325                        command += " --target={}".format(targets)
326                        build_commands.append(command)
327
328        return build_commands
329
330    # @Override
331    def SetUp(self):
332        """Initializes the parser for config command."""
333        self.schedule_thread = {}
334        self.arg_parser.add_argument(
335            "--update",
336            choices=("single", "start", "stop", "list"),
337            default="start",
338            help="Update build info")
339        self.arg_parser.add_argument(
340            "--id",
341            default=None,
342            help="session ID only required for 'stop' update command")
343        self.arg_parser.add_argument(
344            "--interval",
345            type=int,
346            default=60,
347            help="Interval (seconds) to repeat build update.")
348        self.arg_parser.add_argument(
349            "--config-type",
350            choices=("prod", "test"),
351            default="prod",
352            help="Whether it's for prod")
353        self.arg_parser.add_argument(
354            "--branch",
355            required=True,
356            help="Branch to grab the artifact from.")
357        self.arg_parser.add_argument(
358            "--target",
359            required=True,
360            help="a comma-separate list of build target product(s).")
361        self.arg_parser.add_argument(
362            "--account_id",
363            default=common._DEFAULT_ACCOUNT_ID,
364            help="Partner Android Build account_id to use.")
365        self.arg_parser.add_argument(
366            '--method',
367            default='GET',
368            choices=('GET', 'POST'),
369            help='Method for fetching')
370        self.arg_parser.add_argument(
371            '--update_build',
372            dest='update_build',
373            action='store_true',
374            help='A boolean value indicating whether to upload build info.')
375        self.arg_parser.add_argument(
376            "--clear_schedule",
377            default=False,
378            help="True to clear all schedule data on the scheduler cloud")
379        self.arg_parser.add_argument(
380            "--clear_labinfo",
381            default=False,
382            help="True to clear all lab info data on the scheduler cloud")
383
384    # @Override
385    def Run(self, arg_line):
386        """Updates global config."""
387        args = self.arg_parser.ParseLine(arg_line)
388        if args.update == "single":
389            self.UpdateConfig(args.account_id, args.branch, args.target,
390                              args.config_type, args.method, args.update_build,
391                              args.clear_schedule, args.clear_labinfo)
392        elif args.update == "list":
393            logging.info("Running config update sessions:")
394            for id in self.schedule_thread:
395                logging.info("  ID %d", id)
396        elif args.update == "start":
397            if args.interval <= 0:
398                raise ConsoleArgumentError("update interval must be positive")
399            # do not allow user to create new
400            # thread if one is currently running
401            if args.id is None:
402                if not self.schedule_thread:
403                    args.id = 1
404                else:
405                    args.id = max(self.schedule_thread) + 1
406            else:
407                args.id = int(args.id)
408            if args.id in self.schedule_thread and not hasattr(
409                    self.schedule_thread[args.id], 'keep_running'):
410                logging.warning('config update already running. '
411                                'run config --update=stop --id=%s first.',
412                                args.id)
413                return
414            self.schedule_thread[args.id] = threading.Thread(
415                target=self.UpdateConfigLoop,
416                args=(
417                    args.account_id,
418                    args.branch,
419                    args.target,
420                    args.config_type,
421                    args.method,
422                    args.update_build,
423                    args.interval,
424                    args.clear_schedule,
425                    args.clear_labinfo,
426                ))
427            self.schedule_thread[args.id].daemon = True
428            self.schedule_thread[args.id].start()
429        elif args.update == "stop":
430            if args.id is None:
431                logging.error("--id must be set for stop")
432            else:
433                self.schedule_thread[int(args.id)].keep_running = False
434
435    def Help(self):
436        base_command_processor.BaseCommandProcessor.Help(self)
437        logging.info("Sample: config --branch=<branch name> "
438                     "--target=<build target> "
439                     "--account_id=<account id> --config-type=[prod|test] "
440                     "--update=single")
441