frc971/downloader/downloader.py

# This file is run by shell scripts generated by the aos_downloader Skylark
# macro. Everything before the first -- is a hard-coded list of files to
# download.

from __future__ import print_function

import argparse
import sys
import subprocess
import re
import os


def install(ssh_target, pkg, ssh_path, scp_path):
    """Installs a package from NI on the ssh target."""
    print("Installing", pkg)
    PKG_URL = "http://download.ni.com/ni-linux-rt/feeds/2015/arm/ipk/cortexa9-vfpv3/" + pkg
    subprocess.check_call(["wget", PKG_URL, "-O", pkg])
    try:
        subprocess.check_call(
            [scp_path, "-S", ssh_path, pkg, ssh_target + ":/tmp/" + pkg])
        subprocess.check_call(
            [ssh_path, ssh_target, "opkg", "install", "/tmp/" + pkg])
        subprocess.check_call([ssh_path, ssh_target, "rm", "/tmp/" + pkg])
    finally:
        subprocess.check_call(["rm", pkg])


def main(argv):
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--target",
        type=str,
        default="roborio-971-frc.local",
        help="Target to deploy code to.")
    parser.add_argument(
        "--type",
        type=str,
        choices=["roborio", "pi"],
        required=True,
        help="Target type for deployment")
    parser.add_argument(
        "--dir",
        type=str,
        help="Directory within robot_code to copy the files to.")
    parser.add_argument(
        "srcs", type=str, nargs='+', help="List of files to copy over")
    args = parser.parse_args(argv[1:])

    relative_dir = ""
    recursive = False

    srcs = args.srcs
    if args.dir is not None:
        relative_dir = args.dir
        recursive = True

    destination = args.target

    result = re.match("(?:([^:@]+)@)?([^:@]+)(?::([^:@]+))?", destination)
    if not result:
        print(
            "Not sure how to parse destination \"%s\"!" % destination,
            file=sys.stderr)
        return 1
    user = None
    if result.group(1):
        user = result.group(1)
    hostname = result.group(2)

    if result.group(3):
        target_dir = result.group(3)

    if user is None:
        if args.type == "pi":
            user = "pi"
        elif args.type == "roborio":
            user = "admin"
    target_dir = "/home/" + user + "/robot_code"

    ssh_target = "%s@%s" % (user, hostname)

    ssh_path = "external/ssh/ssh"
    scp_path = "external/ssh/scp"

    rsync_cmd = ([
        "external/rsync/usr/bin/rsync", "-e", ssh_path, "-c", "-v", "-z",
        "--perms", "--copy-links"
    ] + srcs)

    # If there is only 1 file to transfer, we would overwrite the destination
    # folder.  In that case, specify the full path to the target.
    if len(srcs) == 1:
        rsync_cmd += [
            "%s:%s/%s/%s" % (ssh_target, target_dir, relative_dir, srcs[0])
        ]
    else:
        rsync_cmd += ["%s:%s/%s" % (ssh_target, target_dir, relative_dir)]

    try:
        subprocess.check_call(rsync_cmd)
    except subprocess.CalledProcessError as e:
        if e.returncode == 127 or e.returncode == 12:
            print("Unconfigured roboRIO, installing rsync.")
            install(ssh_target, "libattr1_2.4.47-r0.36_cortexa9-vfpv3.ipk",
                    ssh_path, scp_path)
            install(ssh_target, "libacl1_2.2.52-r0.36_cortexa9-vfpv3.ipk",
                    ssh_path, scp_path)
            install(ssh_target, "rsync_3.1.0-r0.7_cortexa9-vfpv3.ipk",
                    ssh_path, scp_path)
            subprocess.check_call(rsync_cmd)
        elif e.returncode == 11:
            # Directory wasn't created, make it and try again.  This keeps the happy path fast.
            subprocess.check_call(
                [ssh_path, ssh_target, "mkdir", "-p", target_dir])
            subprocess.check_call(rsync_cmd)
        else:
            raise e

    if not recursive:
        subprocess.check_call((ssh_path, ssh_target, "&&".join([
            "chmod u+s %s/starterd.stripped" % target_dir,
            "echo \'Done moving new executables into place\'",
            "bash -c \'sync\'",
        ])))


if __name__ == "__main__":
    main(sys.argv)