go/chromeos/setup_chromeos_testing.py

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

"""Generate board-specific scripts for Go compiler testing."""


import argparse
import getpass
import os
import sys

from cros_utils import command_executer

SUCCESS = 0
DEBUG = False

ARCH_DATA = {"x86_64": "amd64", "arm32": "arm", "arm64": "arm64"}

CROS_TOOLCHAIN_DATA = {
    "x86_64": "x86_64-cros-linux-gnu",
    "arm32": "armv7a-cros-linux-gnueabihf",
    "arm64": "aarch64-cros-linux-gnu",
}

GLIBC_DATA = {"x86_64": "glibc", "arm32": "glibc32", "arm64": "glibc"}

CONFIG_TEMPLATE = """
Host %s
    HostName %s
    User root
    UserKnownHostsFile /dev/null
    BatchMode yes
    CheckHostIP no
    StrictHostKeyChecking no
    IdentityFile %%d/.ssh/testing_rsa
"""

BASE_TEMPLATE = """#!/bin/bash

# Invoke the Go cross compiler for %s.
# Uses ../go_target to add PIE flags.
#
# This is just an example for an %s board.

GOOS="linux" GOARCH="%s" CGO_ENABLED="1" \\
        GOROOT="/usr/lib/go/%s" \\
        CC="%s-clang" \\
        CXX="%s-clang++" \\
        exec go_target "$@"
"""

EXEC_TEMPLATE = """#!/bin/bash

# Copy and remotely execute a binary on the %s device.
#
# For this to work, the corresponding entry must exist in
# ~/.ssh/config and the device must already be setup for
# password-less authentication. See setup instructions at
# http://go/chromeos-toolchain-team/go-toolchain

GOOS="linux" GOARCH="%s" \\
      GOLOADER="/tmp/%s/ld.so" \\
      exec go_target_exec %s "$@"
"""


def log(msg):

    if DEBUG:
        print(msg)


def WriteFile(file_content, file_name):
    with open(file_name, "w", encoding="utf-8") as out_file:
        out_file.write(file_content)


def GenerateGoHelperScripts(ce, arm_board, x86_board, chromeos_root):
    keys = ["x86_64", "arm32", "arm64"]
    names = {
        "x86_64": x86_board,
        "arm64": arm_board,
        "arm32": ("%s32" % arm_board),
    }

    toolchain_dir = os.path.join(
        chromeos_root, "src", "third_party", "toolchain-utils", "go", "chromeos"
    )
    for k in keys:
        name = names[k]
        arch = ARCH_DATA[k]
        toolchain = CROS_TOOLCHAIN_DATA[k]
        glibc = GLIBC_DATA[k]

        base_file = os.path.join(toolchain_dir, ("go_%s" % name))
        base_file_content = BASE_TEMPLATE % (
            name,
            arch,
            arch,
            toolchain,
            toolchain,
            toolchain,
        )
        WriteFile(base_file_content, base_file)
        cmd = "chmod 755 %s" % base_file
        ce.RunCommand(cmd)

        exec_file = os.path.join(toolchain_dir, ("go_%s_exec" % name))
        exec_file_content = EXEC_TEMPLATE % (name, arch, glibc, name)
        WriteFile(exec_file_content, exec_file)
        cmd = "chmod 755 %s" % exec_file
        ce.RunCommand(cmd)

    return 0


def UpdateChrootSshConfig(
    ce, arm_board, arm_dut, x86_board, x86_dut, chromeos_root
):
    log("Entering UpdateChrootSshConfig")
    # Copy testing_rsa to .ssh and set file protections properly.
    user = getpass.getuser()
    ssh_dir = os.path.join(chromeos_root, "chroot", "home", user, ".ssh")
    dest_file = os.path.join(ssh_dir, "testing_rsa")
    src_file = os.path.join(
        chromeos_root, "src", "scripts", "mod_for_test_scripts", "testing_rsa"
    )
    if not os.path.exists(dest_file):
        if os.path.exists(src_file):
            cmd = "cp %s %s" % (src_file, dest_file)
            ret = ce.RunCommand(cmd)
            if ret != SUCCESS:
                print('Error executing "%s". Exiting now...' % cmd)
                sys.exit(1)
            cmd = "chmod 600 %s" % dest_file
            ret = ce.RunCommand(cmd)
            if ret != SUCCESS:
                print(
                    "Error executing %s; may need to re-run this manually."
                    % cmd
                )
        else:
            print(
                "Cannot find %s; you will need to update testing_rsa by hand."
                % src_file
            )
    else:
        log("testing_rsa exists already.")

    # Save ~/.ssh/config file, if not already done.
    config_file = os.path.expanduser("~/.ssh/config")
    saved_config_file = os.path.join(
        os.path.expanduser("~/.ssh"), "config.save.go-scripts"
    )
    if not os.path.exists(saved_config_file):
        cmd = "cp %s %s" % (config_file, saved_config_file)
        ret = ce.RunCommand(cmd)
        if ret != SUCCESS:
            print("Error making save copy of ~/.ssh/config. Exiting...")
            sys.exit(1)

    # Update ~/.ssh/config file
    log("Reading ssh config file")
    with open(config_file, "r") as input_file:
        config_lines = input_file.read()

    x86_host_config = CONFIG_TEMPLATE % (x86_board, x86_dut)
    arm_names = "%s %s32" % (arm_board, arm_board)
    arm_host_config = CONFIG_TEMPLATE % (arm_names, arm_dut)

    config_lines += x86_host_config
    config_lines += arm_host_config

    log("Writing ~/.ssh/config")
    WriteFile(config_lines, config_file)

    return 0


def CleanUp(ce, x86_board, arm_board, chromeos_root):
    # Find and remove go helper scripts
    keys = ["x86_64", "arm32", "arm64"]
    names = {
        "x86_64": x86_board,
        "arm64": arm_board,
        "arm32": ("%s32" % arm_board),
    }

    toolchain_dir = os.path.join(
        chromeos_root, "src", "third_party", "toolchain-utils", "go", "chromeos"
    )
    for k in keys:
        name = names[k]
        base_file = os.path.join(toolchain_dir, ("go_%s" % name))
        exec_file = os.path.join(toolchain_dir, ("go_%s_exec" % name))
        cmd = "rm -f %s; rm -f %s" % (base_file, exec_file)
        ce.RunCommand(cmd)

    # Restore saved config_file
    config_file = os.path.expanduser("~/.ssh/config")
    saved_config_file = os.path.join(
        os.path.expanduser("~/.ssh"), "config.save.go-scripts"
    )
    if not os.path.exists(saved_config_file):
        print(
            "Could not find file: %s; unable to restore ~/.ssh/config ."
            % saved_config_file
        )
    else:
        cmd = "mv %s %s" % (saved_config_file, config_file)
        ce.RunCommand(cmd)

    return 0


def Main(argv):
    # pylint: disable=global-statement
    global DEBUG

    parser = argparse.ArgumentParser()
    parser.add_argument("-a", "--arm64_board", dest="arm_board", required=True)
    parser.add_argument(
        "-b", "--x86_64_board", dest="x86_64_board", required=True
    )
    parser.add_argument(
        "-c", "--chromeos_root", dest="chromeos_root", required=True
    )
    parser.add_argument("-x", "--x86_64_dut", dest="x86_64_dut", required=True)
    parser.add_argument("-y", "--arm64_dut", dest="arm_dut", required=True)
    parser.add_argument(
        "-z", "--cleanup", dest="cleanup", default=False, action="store_true"
    )
    parser.add_argument(
        "-v", "--verbose", dest="verbose", default=False, action="store_true"
    )

    options = parser.parse_args(argv[1:])

    if options.verbose:
        DEBUG = True

    if not os.path.exists(options.chromeos_root):
        print("Invalid ChromeOS Root: %s" % options.chromeos_root)

    ce = command_executer.GetCommandExecuter()
    all_good = True
    for m in (options.x86_64_dut, options.arm_dut):
        cmd = "ping -c 3 %s > /dev/null" % m
        ret = ce.RunCommand(cmd)
        if ret != SUCCESS:
            print("Machine %s is currently not responding to ping." % m)
            all_good = False

    if not all_good:
        return 1

    if not options.cleanup:
        UpdateChrootSshConfig(
            ce,
            options.arm_board,
            options.arm_dut,
            options.x86_64_board,
            options.x86_64_dut,
            options.chromeos_root,
        )
        GenerateGoHelperScripts(
            ce, options.arm_board, options.x86_64_board, options.chromeos_root
        )
    else:
        CleanUp(
            ce, options.x86_64_board, options.arm_board, options.chromeos_root
        )

    return 0


if __name__ == "__main__":
    val = Main(sys.argv)
    sys.exit(val)