#!/usr/bin/env python3
# Copyright (c) 2015 Arista Networks, Inc.  All rights reserved.
# Arista Networks, Inc. Confidential and Proprietary.

import Tac, Tracing, QuickTrace
import subprocess, argparse, sys, collections, os.path

traceDebug = Tracing.trace0
qv = QuickTrace.Var
qtDebug = QuickTrace.trace0

RemoteHost = collections.namedtuple( 'RemoteHost', [ 'host', 'port',
                                                     'user', 'key' ] )
errors = 0

def runCmd( cmd ): # pylint: disable=useless-return
   try:
      traceDebug( 'Subprocess Cmd:', cmd )
      subprocess.check_output( cmd )
   except subprocess.CalledProcessError as err:
      qtDebug( 'Cmd: ', cmd, ' failed with return code: ', err.returncode )
      traceDebug( 'Cmd: ', cmd, ' failed with return code: ', err.returncode )
      global errors
      errors += 1
   return

def _sshOptions( remoteHost ):
   # pylint: disable-next=consider-using-f-string
   options = [ '-oPort=%s' % remoteHost.port,
               # pylint: disable-next=consider-using-f-string
               '-oUser=%s' % remoteHost.user,
               '-oStrictHostKeyChecking=no',
               '-oUserKnownHostsFile=/dev/null',
               '-oCheckHostIP=no',
               '-oLogLevel=quiet' ]
   if remoteHost.key:
      # pylint: disable-next=consider-using-f-string
      options.append( '-oIdentityFile=%s' % remoteHost.key )
   return options

def sshCommand( cmd, remoteHost ):
   remoteCmd = [ 'ssh' ]
   remoteCmd += _sshOptions( remoteHost )
   remoteCmd += [ f'{remoteHost.user}@{remoteHost.host}' ]
   return remoteCmd + cmd

def createDir( target, remoteHost=None ):
   targetDir = os.path.dirname( target )
   if not targetDir:
      return
   mkdirCmd = [ 'mkdir', '-p', targetDir ]

   if remoteHost:
      cmd = sshCommand( mkdirCmd, remoteHost ) 
      traceDebug( 'Remote mkdir:', cmd )
      runCmd( cmd )
   else:
      traceDebug( 'Local mkdir command:', mkdirCmd )
      runCmd( mkdirCmd )

def rsyncFile( remoteHosts, source, target, rsyncOptions ):
   traceDebug( 'rsyncFile:', remoteHosts, source, target, rsyncOptions )
   # pylint: disable-next=consider-using-f-string
   rsyncOptions = [ '--%s' % option for option in rsyncOptions.split() ]
   traceDebug( 'rsync:', rsyncOptions )
   if not remoteHosts:
      createDir( target )
      cmd = [ 'rsync', '-avI', '--no-p', '--delete' ] + rsyncOptions
      cmd += [ source, target ]
      traceDebug( 'Local rsyncFile cmd:', cmd )
      runCmd( cmd )
   else:
      # rsync to one remoteHost after another
      for r in remoteHosts:
         remoteHost = RemoteHost( *r.split( ':' )[ : 4 ] )
         createDir( target, remoteHost )
         cmd = [ 'rsync', '-avI', '--no-p', '--delete' ] + rsyncOptions
         cmd += [ '--rsh',
                  # pylint: disable-next=consider-using-f-string
                  '/usr/bin/ssh %s' % ' '.join( _sshOptions( remoteHost ) ),
                  source, f'{remoteHost.host}:{target}' ]
         traceDebug( 'rsyncFile cmd:', cmd )
         runCmd( cmd )
   sys.exit( errors )

def deleteFile( remoteHosts, target ):
   traceDebug( 'deleteFile:', remoteHosts, target )
   rmCmd = [ 'rm', '-rf', target ]
   if not remoteHosts:
      traceDebug( 'Local rm cmd:', rmCmd )
      runCmd( rmCmd )
   else:
      for r in remoteHosts:
         remoteHost = RemoteHost( *r.split( ':' )[ : 4 ] )
         cmd = sshCommand( rmCmd, remoteHost )
         traceDebug( 'deleteFile: cmd=', cmd )
         runCmd( cmd )
   sys.exit( errors )

def parseArguments():
   parser = argparse.ArgumentParser()
   parser.add_argument( '--cmd' )
   parser.add_argument( '--remoteHosts' )
   parser.add_argument( '--source' )
   parser.add_argument( '--target' )
   parser.add_argument( '--rsyncOptions' )
   return parser.parse_args()

def main():
   traceDebug( 'sys.argv = ', sys.argv )
   arguments = parseArguments()
   traceDebug( 'main: arguments=', arguments )
   if arguments.cmd == 'rsyncFile':
      if arguments.remoteHosts:
         arguments.remoteHosts = arguments.remoteHosts.split()
      rsyncFile( arguments.remoteHosts, arguments.source, arguments.target,
                 arguments.rsyncOptions )
   elif arguments.cmd == 'deleteFile':
      if arguments.remoteHosts:
         arguments.remoteHosts = arguments.remoteHosts.split()
      deleteFile( arguments.remoteHosts, arguments.target )
   else:
      assert False

if __name__ == '__main__':
   main()
