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)