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

import os
import pty
import sys

args = sys.argv
if len( args ) < 3:
   print( "Error: usage: aScpPass <password> <cmdline>" )
   sys.exit( -1 )

password = args[1]
execCmd = args[2:]

# Run "cmd", then wait for "needle" to show up in "fd".
# cmd can be None, say after initial connection, synching for a prompt (needle).
# needle can be None, when we no longer caring about further output (like progress)
def runwait( fd_, cmd, needle ):
   try:
      if cmd:
         os.write( fd_, cmd.encode( 'utf-8' ) )
      if needle:
         buf_ = ""
         while 1:
            buf_ = buf_ + os.read( fd_, 1024 ).decode( 'utf-8' )
            if needle in buf_:
               return ( True, buf_ )
   except OSError:
      return ( False, buf_ )
   return ( False, "" ) # make pylint happy

# Create a terminal to run scp in (scp wants to be interactive -> terminal needed)
# Fork and make the child a session leader with a controlling terminal with "fd"
# opened in parent and connected to the child's input.
pid, fd = pty.fork()
if pid == 0:
   # Start scp program in child
   os.execvp( "scp", execCmd )
else:
   # Monitor output of scp from parent, feed password when prompted, then display
   # the progress bar as it happens (or eventual error).
   # Wait for password prompt
   res = runwait( fd, None, "assword" )
   if res[ 0 ] is False: # program terminated before getting password prompt
      status = os.waitpid( pid, 0 )
      if status[ 1 ] == 0: # password was not needed (and it worked anyway)
         sys.exit( 0 )
      print( res[ 1 ] ) # "Could not resolve hostname" or some such error
      sys.exit( -1 ) # status[1] is usually 256, which is 0...
   # Send password
   res = runwait( fd, password + "\n", None )
   # Print the prompt (or maybe some "No such file or directory" error)
   # Read byte by byte since this is a progress bar (no <cr> for a while)
   try:
      buf = "" # collect output to eventully check if file was 100% transfered
      while True:
         c = os.read( fd, 1 ).decode( 'utf-8' )
         if not buf and c in [ '\r', '\n' ]: # first we see <cr> from the password
            continue
         print( c, end='' )
         buf = buf + c
   except OSError:
      pass
   # Determine success or error
   status = os.waitpid( pid, 0 )
   if status[ 1 ] == 0 and "100%" in buf:
      sys.exit( 0 )
   sys.exit(-2)

