path: root/contrib/upload_results.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
# Copyright 2020 The Chromium OS Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

import argparse
import logging as log
import os
import re
import shlex
import shutil
import subprocess
import multiprocessing
import sys
import time
import uuid
import json
import functools
import glob

from google.cloud import storage
from google.api_core import exceptions as cloud_exceptions
# pylint: disable=no-name-in-module, import-error

import common
from autotest_lib.client.common_lib import global_config
from autotest_lib.client.common_lib import mail, pidfile
from autotest_lib.tko.parse import parse_one, export_tko_job_to_file

# Appends the moblab source paths for the pubsub wrapper
sys.path.append('/mnt/host/source/src/platform/moblab/src')
from moblab_common import pubsub_client

STATUS_FILE = "status"
STATUS_LOG_FILE = "status.log"
KEYVAL_FILE = "keyval"
NEW_KEYVAL_FILE = "new_keyval"
UPLOADED_STATUS_FILE = ".uploader_status"
STATUS_GOOD = "PUBSUB_SENT"
FAKE_MOBLAB_ID_FILE = "fake_moblab_id_do_not_delete.txt"
GIT_HASH_FILE = "git_hash.txt"
GIT_COMMAND = ("git log --pretty=format:'%h -%d %s (%ci) <%an>'"
               " --abbrev-commit -20")
AUTOTEST_DIR = "/mnt/host/source/src/third_party/autotest/files/"
DEFAULT_SUITE_NAME = "default_suite"
SUITE_NAME_REGEX = "Fetching suite for suite named (.+?)\.\.\."
DEBUG_FILE_PATH = "debug/test_that.DEBUG"
CONFIG_DIR = os.path.dirname(os.path.abspath(__file__)) + "/config/"
DEFAULT_BOTO_CONFIG = CONFIG_DIR + ".boto_upload_utils"
UPLOAD_CONFIG = CONFIG_DIR + "upload_config.json"
SERVICE_ACCOUNT_CONFIG = CONFIG_DIR + ".service_account.json"

logging = log.getLogger(__name__)


def parse_arguments(argv):
    """Creates the argument parser.

    Args:
        argv: A list of input arguments.

    Returns:
        A parser object for input arguments.
    """
    parser = argparse.ArgumentParser(description=__doc__)
    subparsers = parser.add_subparsers(
            help='select sub option for test result utility',
            dest='subcommand')
    subparsers.required = True
    parser.add_argument("-v",
                        "--verbose",
                        dest='verbose',
                        action='store_true',
                        help="Enable verbose (debug) logging.")
    parser.add_argument("-q",
                        "--quiet",
                        dest='quiet',
                        action='store_true',
                        help="Quiet mode for background call")
    def_logfile = "/tmp/" + os.path.basename(
            sys.argv[0].split(".")[0]) + ".log"
    parser.add_argument("-l",
                        "--logfile",
                        type=str,
                        required=False,
                        default=def_logfile,
                        help="Full path to logfile. Default: " + def_logfile)

    # configuration subcommand to create config file and populate environment
    config_parser = subparsers.add_parser(name="config",
                                          help='upload test results to CPCon')
    config_parser.add_argument(
            "-b",
            "--bucket",
            type=str,
            required=True,
            help="The GCS bucket that test results are uploaded to, e.g."
            "'gs://xxxx'.")
    config_parser.add_argument("-f",
                               "--force",
                               dest='force',
                               action="store_true",
                               help="Force overwrite of previous config files")

    upload_parser = subparsers.add_parser(name="upload",
                                          help='upload test results to CPCon')
    upload_parser.add_argument(
            "--bug",
            type=_valid_bug_id,
            required=False,
            help=
            "Write bug id to the test results. Each test entry can only have "
            "at most 1 bug id. Optional.")
    upload_parser.add_argument(
            "-d",
            "--directory",
            type=str,
            required=True,
            help="The directory of non-Moblab test results.")
    upload_parser.add_argument(
            "--parse_only",
            action='store_true',
            help="Generate job.serialize locally but do not upload test "
            "directories and not send pubsub messages.")
    upload_parser.add_argument(
            "--upload_only",
            action='store_true',
            help="Leave existing protobuf files as-is, only upload "
            "directories and send pubsub messages.")
    upload_parser.add_argument(
            "-f",
            "--force",
            dest='force',
            action='store_true',
            help=
            "force re-upload of results even if results were already successfully uploaded."
    )
    upload_parser.add_argument(
            "-s",
            "--suite",
            type=str,
            default=None,
            help="The suite is used to identify the type of test results,"
            "e.g. 'power' for platform power team. If not specific, the "
            "default value is 'default_suite'.")
    return parser.parse_args(argv)


def _confirm_option(question):
    """
        Get a yes/no answer from the user via command line.

    Args:
        question: string, question to ask the user.

    Returns:
        A boolean. True if yes; False if no.
    """
    expected_answers = ['y', 'yes', 'n', 'no']
    answer = ''
    while answer not in expected_answers:
        answer = input(question + "(y/n): ").lower().strip()
    return answer[0] == "y"


def _read_until_string(pipe, stop_string):
    lines = [""]
    while True:
        c = pipe.read(1)
        lines[-1] = lines[-1] + c.decode("utf-8")
        if stop_string == lines[-1]:
            return lines
        if c.decode("utf-8") == "\n":
            lines.append("")


def _configure_environment(parsed_args):
    # create config directory if not exists
    os.makedirs(CONFIG_DIR, exist_ok=True)

    if os.path.exists(UPLOAD_CONFIG) and not parsed_args.force:
        logging.error("Environment already configured, run with --force")
        exit(1)

    # call the gsutil config tool to set up accounts
    if os.path.exists(DEFAULT_BOTO_CONFIG + ".bak"):
        os.remove(DEFAULT_BOTO_CONFIG + ".bak")

    if os.path.exists(DEFAULT_BOTO_CONFIG):
        os.remove(DEFAULT_BOTO_CONFIG)
    os.mknod(DEFAULT_BOTO_CONFIG)
    os.environ["BOTO_CONFIG"] = DEFAULT_BOTO_CONFIG
    os.environ[
            "GOOGLE_APPLICATION_CREDENTIALS"] = CONFIG_DIR + ".service_account.json"
    with subprocess.Popen(["gsutil", "config"],
                          stdout=subprocess.PIPE,
                          stderr=subprocess.PIPE,
                          stdin=subprocess.PIPE) as sp:
        lines = _read_until_string(sp.stdout, "Enter the authorization code: ")
        code = input("enter auth code from " + str(lines[1]) + ": ")
        sp.stdin.write(bytes(code + '\n', "utf-8"))
        sp.stdin.flush()
        lines = _read_until_string(sp.stdout, "What is your project-id? ")
        sp.stdin.write(bytes(parsed_args.bucket + '\n', "utf-8"))
        sp.stdin.flush()

    subprocess.run([
            "gsutil", "cp",
            "gs://" + parsed_args.bucket + "/.service_account.json", CONFIG_DIR
    ])
    subprocess.run([
            "gsutil", "cp",
            "gs://" + parsed_args.bucket + "/pubsub-key-do-not-delete.json",
            CONFIG_DIR
    ])

    sa_filename = ""
    if os.path.exists(CONFIG_DIR + "/.service_account.json"):
        sa_filename = ".service_account.json"
    elif os.path.exists(CONFIG_DIR + "/pubsub-key-do-not-delete.json"):
        sa_filename = "pubsub-key-do-not-delete.json"
    else:
        logging.error("No pubsub key found in bucket, failed config!")
        exit(1)

    # deposit parsed_args.bucket to the json file
    with open(UPLOAD_CONFIG, "w") as cf:
        settings = {}
        settings["bucket"] = parsed_args.bucket
        settings["service_account"] = CONFIG_DIR + sa_filename
        settings["boto_key"] = DEFAULT_BOTO_CONFIG

        cf.write(json.dumps(settings))


def _load_config():
    mandatory_keys = ["bucket", "service_account", "boto_key"]

    if not os.path.exists(UPLOAD_CONFIG):
        logging.error("Missing mandatory config file, run config command")
        exit(1)
    with open(UPLOAD_CONFIG, "r") as cf:
        settings = json.load(cf)

    for key in mandatory_keys:
        if key not in settings:
            logging.error("Missing mandatory setting " + str(key) +
                          ", run config command")
            exit()

    os.environ["GOOGLE_APPLICATION_CREDENTIALS"] = settings["service_account"]
    os.environ["BOTO_CONFIG"] = settings["boto_key"]
    return argparse.Namespace(**settings)


class ResultsManager:
    def __init__(self, results_parser, results_sender):
        self.parent_directories = []
        self.result_directories = set()
        self.results = []
        self.results_parser = results_parser
        self.results_sender = results_sender
        self.bug_id = None
        self.suite_name = ""

        self.moblab_id = self.get_fake_moblab_id()

    def new_directory(self, parent_dir: str):
        self.parent_directories.append(parent_dir)

    def enumerate_all_directories(self):
        self.result_directories = set()
        for parent_dir in self.parent_directories:
            self.enumerate_result_directories(parent_dir)

    def enumerate_result_directories(self, parent_dir):
        """ Gets all test directories.

        Args:
        parent_dir: The parent directory of one or multiple test directories

        Creates a local_result for all directories with a status.log file
        and appends to local_results
        """
        if not os.path.exists(parent_dir) or not os.path.isdir(parent_dir):
            logging.warning('Test directory does not exist: %s' % parent_dir)
            return

        status_log_file = os.path.join(parent_dir, STATUS_LOG_FILE)
        if os.path.exists(status_log_file):
            self.result_directories.add(parent_dir)
            return

        for dir_name in os.listdir(parent_dir):
            subdir = os.path.join(parent_dir, dir_name)
            if os.path.isdir(subdir):
                self.enumerate_result_directories(subdir)

    def set_destination(self, destination):
        self.results_sender.set_destination(destination)

    def get_fake_moblab_id(self):
        """Get or generate a fake moblab id.

        Moblab id is the unique id to a moblab device. Since the upload script runs
        from the chroot instead of a moblab device, we need to generate a fake
        moblab id to comply with the CPCon backend. If there is a previously saved
        fake moblab id, read and use it. Otherwise, generate a uuid to fake a moblab
        device, and store it in the same directory as the upload script.

        Returns:
            A string representing a fake moblab id.
        """
        script_dir = os.path.dirname(__file__)
        fake_moblab_id_path = os.path.join(script_dir, "config",
                                           FAKE_MOBLAB_ID_FILE)

        # Migrate from prior moblab ID location into config directory if possible
        old_moblab_id_file = os.path.join(script_dir, FAKE_MOBLAB_ID_FILE)
        if os.path.exists(old_moblab_id_file):
            logging.info(
                    'Found an existing moblab ID outside config directory, migrating now'
            )
            os.rename(old_moblab_id_file, fake_moblab_id_path)
        try:
            with open(fake_moblab_id_path, "r") as fake_moblab_id_file:
                fake_moblab_id = str(fake_moblab_id_file.read())[0:32]
                if fake_moblab_id:
                    return fake_moblab_id
        except IOError as e:
            logging.info(
                    'Cannot find a fake moblab id at %s, creating a new one.',
                    fake_moblab_id_path)
        fake_moblab_id = uuid.uuid4().hex
        try:
            with open(fake_moblab_id_path, "w") as fake_moblab_id_file:
                fake_moblab_id_file.write(fake_moblab_id)
        except IOError as e:
            logging.warning('Unable to write the fake moblab id to %s: %s',
                            fake_moblab_id_path, e)
        return fake_moblab_id

    def overwrite_suite_name(self, suite_name):
        self.suite_name = suite_name

    def annotate_results_with_bugid(self, bug_id):
        self.bug_id = bug_id

    def parse_all_results(self, upload_only: bool = False):
        self.results = []
        self.enumerate_all_directories()

        for result_dir in self.result_directories:
            if self.bug_id is not None:
                self.results_parser.write_bug_id(result_dir, self.bug_id)
            self.results.append(
                    (result_dir,
                     self.results_parser.parse(result_dir,
                                               upload_only,
                                               suite_name=self.suite_name)))

    def upload_all_results(self, force):
        for result in self.results:
            self.results_sender.upload_result_and_notify(
                    result[0], self.moblab_id, result[1], force)


class FakeTkoDb:
    def find_job(self, tag):
        return None

    def run_with_retry(self, fn, *args):
        fn(*args)


class ResultsParserClass:
    def __init__(self):
        pass

    def job_tag(self, job_id, machine):
        return str(job_id) + "-moblab/" + str(machine)

    def parse(self, path, upload_only: bool, suite_name=""):
        #temporarily assign a fake job id until parsed
        fake_job_id = 1234
        fake_machine = "localhost"
        name = self.job_tag(fake_job_id, fake_machine)
        parse_options = argparse.Namespace(
                **{
                        "suite_report": False,
                        "dry_run": True,
                        "reparse": False,
                        "mail_on_failure": False
                })
        pid_file_manager = pidfile.PidFileManager("parser", path)
        self.print_autotest_git_history(path)
        job = parse_one(FakeTkoDb(), pid_file_manager, name, path,
                        parse_options)
        job.board = job.tests[0].attributes['host-board']
        job_id = int(job.started_time.timestamp() * 1000)
        job.afe_parent_job_id = job_id + 1
        if suite_name == "":
            job.suite = self.parse_suite_name(path)
        else:
            job.suite = suite_name
        job.build_version = self.get_build_version(job.tests)
        name = self.job_tag(job_id, job.machine)
        if not upload_only:
            export_tko_job_to_file(job, name, path + "/job.serialize")

        # autotest_lib appends additional global logger handlers
        # remove these handlers to avoid affecting logging for the google
        # storage library
        for handler in log.getLogger().handlers:
            log.getLogger().removeHandler(handler)
        return job

    def print_autotest_git_history(self, path):
        """
        Print the hash of the latest git commit of the autotest directory.

        Args:
            path: The test directory for non-moblab test results.
        """
        git_hash = subprocess.check_output(shlex.split(GIT_COMMAND),
                                           cwd=AUTOTEST_DIR)
        git_hash_path = os.path.join(path, GIT_HASH_FILE)
        with open(git_hash_path, "w") as git_hash_file:
            git_hash_file.write(git_hash.decode("utf-8"))

    def parse_suite_name(self, path):
        """Get the suite name from a results directory.

        If we don't find the suite name in the first ten lines of test_that.DEBUG
        then return None.

        Args:
            path: The directory specified on the command line.
        """
        path = path.split('/')[:-1]
        path = '/'.join(path)

        debug_file = os.path.join(path, DEBUG_FILE_PATH)
        if not os.path.exists(debug_file) or not os.path.isfile(debug_file):
            return None
        exp = re.compile(SUITE_NAME_REGEX)
        try:
            with open(debug_file) as f:
                line_count = 0
                for line in f:
                    line_count += 1
                    if line_count > 10:
                        break
                    result = exp.search(line)
                    if not result:
                        continue
                    else:
                        return result.group(1)
        except IOError as e:
            logging.warning('Error trying to read test_that.DEBUG: %s', e)
        return DEFAULT_SUITE_NAME

    def get_build_version(self, tests):
        release_version_label = "CHROMEOS_RELEASE_VERSION"
        milestone_label = "CHROMEOS_RELEASE_CHROME_MILESTONE"
        for test in tests:
            if not test.subdir:
                continue

            release = None
            milestone = None
            if release_version_label in test.attributes:
                release = test.attributes[release_version_label]
            if milestone_label in test.attributes:
                milestone = test.attributes[milestone_label]
            if release and milestone:
                return "R%s-%s" % (milestone, release)

        return ""

    def valid_bug_id(self, v):
        """Check if user input bug id is in valid format.

        Args:
            v: User input bug id in string.
        Returns:
            An int representing the bug id.
        Raises:
            argparse.ArgumentTypeError: if user input bug id has wrong format.
        """
        try:
            bug_id = int(v)
        except ValueError as e:
            raise argparse.ArgumentTypeError(
                    "Bug id %s is not a positive integer: "
                    "%s" % (v, e))
        if bug_id <= 0:
            raise argparse.ArgumentTypeError(
                    "Bug id %s is not a positive integer" % v)
        return bug_id

    def write_bug_id(self, test_dir, bug_id):
        """
            Write the bug id to the test results.

        Args:
            test_dir: The test directory for non-moblab test results.
            bug_id: The bug id to write to the test results.
        Returns:
            A boolean. True if the bug id is written successfully or is the same as
            the old bug id already in test results; False if failed to write the
            bug id, or if the user decides not to overwrite the old bug id already
            in test results.
        """
        old_bug_id = None
        new_keyval = list()

        keyval_file = os.path.join(test_dir, KEYVAL_FILE)
        try:
            with open(keyval_file, 'r') as keyval_raw:
                for line in keyval_raw.readlines():
                    match = re.match(r'bug_id=(\d+)', line)
                    if match:
                        old_bug_id = self.valid_bug_id(match.group(1))
                    else:
                        new_keyval.append(line)
        except IOError as e:
            logging.error(
                    'Cannot read keyval file from %s, skip writing the bug '
                    'id %s: %s', test_dir, bug_id, e)
            return False

        if old_bug_id:
            if old_bug_id == bug_id:
                return True
            overwrite_bug_id = _confirm_option(
                    'Would you like to overwrite bug id '
                    '%s with new bug id %s?' % (old_bug_id, bug_id))
            if not overwrite_bug_id:
                return False

        new_keyval.append('bug_id=%s' % bug_id)
        new_keyval_file = os.path.join(test_dir, NEW_KEYVAL_FILE)
        try:
            with open(new_keyval_file, 'w') as new_keyval_raw:
                for line in new_keyval:
                    new_keyval_raw.write(line)
                new_keyval_raw.write('\n')
            shutil.move(new_keyval_file, keyval_file)
            return True
        except Exception as e:
            logging.error(
                    'Cannot write bug id to keyval file in %s, skip writing '
                    'the bug id %s: %s', test_dir, bug_id, e)
            return False


ResultsParser = ResultsParserClass()
_valid_bug_id = functools.partial(ResultsParserClass.valid_bug_id,
                                  ResultsParser)


class ResultsSenderClass:
    def __init__(self):
        self.gcs_bucket = ""

    def set_destination(self, destination):
        self.gcs_bucket = destination

    def upload_result_and_notify(self, test_dir, moblab_id, job, force):
        job_id = str(int(job.started_time.timestamp() * 1000))
        if self.uploaded(test_dir) and not force:
            return
        self.upload_result(test_dir, moblab_id, job_id, job.machine)
        self.send_pubsub_message(test_dir, moblab_id, job_id)

    def upload_batch_files(self, gs_path, test_dir, files):
        for file in files:
            if not os.path.isfile(file):
                continue
            gs_client_bucket = storage.Client().bucket(self.gcs_bucket)
            # remove trailing slash to ensure dest_file path gets created properly
            test_dir = test_dir.rstrip('/')
            dest_file = gs_path + file.replace(test_dir, "", 1)
            logging.info("uploading file: %s", dest_file)
            blob = gs_client_bucket.blob(dest_file)
            blob.upload_from_filename(file)

    def upload_result(self, test_dir, moblab_id, job_id, hostname):
        """
            Upload the test directory with job.serialize to GCS bucket.

        Args:
            args: A list of input arguments.
            test_dir: The test directory for non-moblab test results.
            job_keyval: The key-value object of the job.
            moblab_id: A string that represents the unique id of a moblab device.
            job_id: A job id.
        """
        upload_status_file = os.path.join(test_dir, UPLOADED_STATUS_FILE)
        with open(upload_status_file, "w") as upload_status:
            upload_status.write("UPLOADING")

        fake_moblab_id = moblab_id
        fake_moblab_install_id = moblab_id

        gcs_bucket_path = os.path.join("results", fake_moblab_id,
                                       fake_moblab_install_id,
                                       "%s-moblab" % job_id, hostname)

        try:
            logging.info(
                    "Start to upload test directory: %s to GCS bucket path: %s",
                    test_dir, gcs_bucket_path)
            with open(upload_status_file, "w") as upload_status:
                upload_status.write("UPLOADED")

            files_to_upload = glob.glob(test_dir + "/**", recursive=True)
            batch_size = 8
            with multiprocessing.Pool(4) as p:
                files_to_upload_batch = [
                        files_to_upload[i:i + batch_size]
                        for i in range(0, len(files_to_upload), batch_size)
                ]
                p.map(
                        functools.partial(
                                ResultsSenderClass.upload_batch_files, self,
                                gcs_bucket_path, test_dir),
                        files_to_upload_batch)

            logging.info(
                    "Successfully uploaded test directory: %s to GCS bucket path: %s",
                    test_dir, gcs_bucket_path)
        except Exception as e:
            with open(upload_status_file, "w") as upload_status:
                upload_status.write("UPLOAD_FAILED")
            raise Exception(
                    "Failed to upload test directory: %s to GCS bucket "
                    "path: %s for the error: %s" %
                    (test_dir, gcs_bucket_path, e))

    def send_pubsub_message(self, test_dir, moblab_id, job_id):
        """
            Send pubsub messages to trigger CPCon pipeline to process non-moblab
            test results in the specific GCS bucket path.

        Args:
            bucket: The GCS bucket.
            moblab_id: A moblab id.
            job_id: A job id.
        """
        moblab_install_id = moblab_id
        console_client = pubsub_client.PubSubBasedClient()
        gsuri = "gs://%s/results/%s/%s/%s-moblab" % (
                self.gcs_bucket, moblab_id, moblab_install_id, job_id)

        try:
            logging.info("Start to send the pubsub message to GCS path: %s",
                         gsuri)
            message_id = \
                console_client.send_test_job_offloaded_message(gsuri,
                                                            moblab_id,
                                                            moblab_install_id)
            upload_status_file = os.path.join(test_dir, UPLOADED_STATUS_FILE)
            with open(upload_status_file, "w") as upload_status:
                upload_status.write(STATUS_GOOD)

            logging.info(
                    "Successfully sent the pubsub message with message id: %s to GCS "
                    "path: %s", message_id[0], gsuri)
        except Exception as e:
            raise Exception(
                    "Failed to send the pubsub message with moblab id: %s "
                    "and job id: %s to GCS path: %s for the error: %s" %
                    (moblab_id, job_id, gsuri, e))

    def uploaded(self, test_dir):
        """
        Checks if the message for the uploaded bucket has been sent.

        Args:
            test_dir: The test directory for non-moblab test results.
        """
        upload_status_file = os.path.join(test_dir, UPLOADED_STATUS_FILE)
        if not os.path.exists(upload_status_file):
            logging.debug("The upload status file %s does not exist.",
                          upload_status_file)
            return False

        with open(upload_status_file, "r") as upload_status:
            if upload_status.read() == STATUS_GOOD:
                logging.warn(
                        "The test directory: %s status has already been "
                        "sent to CPCon and the .upload_status file has "
                        "been set to PUBSUB_SENT.", test_dir)
                return True
            else:
                logging.debug("The pubsub message was not successful")
        return False


ResultsSender = ResultsSenderClass()


def main(args):
    parsed_args = parse_arguments(args)

    fmt = log.Formatter('%(asctime)s :: %(levelname)-8s :: %(message)s')
    logging.propagate = False

    log_level = log.INFO
    if parsed_args.verbose:
        log_level = log.DEBUG
    if not parsed_args.quiet:
        stream_handler = log.StreamHandler(sys.stdout)
        stream_handler.setFormatter(fmt)
        stream_handler.setLevel(log_level)
        logging.addHandler(stream_handler)

    logging.info("logging to %s", parsed_args.logfile)
    file_handler = log.FileHandler(parsed_args.logfile, mode='w')
    file_handler.setFormatter(fmt)
    file_handler.setLevel(log.DEBUG)
    logging.addHandler(file_handler)

    if parsed_args.subcommand == "config":
        _configure_environment(parsed_args)
        return

    persistent_settings = _load_config()

    results_manager = ResultsManager(ResultsParser, ResultsSender)
    results_manager.set_destination(persistent_settings.bucket)
    results_manager.new_directory(parsed_args.directory)

    if parsed_args.bug:
        results_manager.annotate_results_with_bugid(parsed_args.bug)
    if parsed_args.suite:
        results_manager.overwrite_suite_name(parsed_args.suite)
    if parsed_args.parse_only:
        results_manager.parse_all_results()
    elif parsed_args.upload_only:
        results_manager.parse_all_results(upload_only=True)
        results_manager.upload_all_results(force=parsed_args.force)
    else:
        results_manager.parse_all_results()
        results_manager.upload_all_results(force=parsed_args.force)


if __name__ == "__main__":
    try:
        main(sys.argv[1:])
    except KeyboardInterrupt:
        sys.exit(0)