aos/vision/download/downloader.py

#!/usr/bin/python3

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

from __future__ import print_function

import sys
import subprocess
import re
import os
import os.path


def RunAndSplitShas(shcmd):
    out = subprocess.check_output(shcmd)
    return [
        line.split(' ')[0] for line in filter(lambda x: x, out.split('\n'))
    ]


def GetChecksums(fnames):
    return RunAndSplitShas(["sha256sum"] + fnames)


def GetJetsonChecksums(ssh_target, fnames):
    target_files = ["/root/%s" % fname for fname in fnames]
    subprocess.check_call(
        ["ssh", ssh_target, "touch " + " ".join(target_files)])
    cmds = ["ssh", ssh_target, "sha256sum " + " ".join(target_files)]
    return RunAndSplitShas(cmds)


def ToJetsonFname(fname):
    if (fname[-9:] == ".stripped"):
        fname = fname[:-9]
    return os.path.basename(fname)


def VerifyCheckSumsAndUpload(fnames, ssh_target):
    jetson_fnames = [ToJetsonFname(fname) for fname in fnames]
    checksums = GetChecksums(fnames)
    jetson_checksums = GetJetsonChecksums(ssh_target, jetson_fnames)
    for i in range(len(fnames)):
        if (checksums[i] != jetson_checksums[i]):
            # if empty, unlink
            subprocess.check_call(
                ["ssh", ssh_target, "unlink " + jetson_fnames[i]])
            subprocess.check_call(
                ["scp", fnames[i], ssh_target + ":" + jetson_fnames[i]])


def main(argv):
    args = argv[argv.index('--') + 1:]
    files = argv[1:argv.index('--')]

    VerifyCheckSumsAndUpload(files, args[-1])


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