#!/usr/bin/python3
import os
import sys
import argparse
import paramiko

def establishSSHConnection( hostname, username, keyFile ):
    # Load the private key file
    serverAccessKey = paramiko.RSAKey.from_private_key(keyFile)
    # Initiate connection to the upload host
    client = paramiko.SSHClient()
    client.set_missing_host_key_policy( paramiko.client.AutoAddPolicy() )
    client.connect( hostname, username=username, pkey = serverAccessKey )
    # All done
    return client

def sftpFileExists(sftpConnection, pathToFile):
    # Try to stat the file
    # If it succeeds, then it exists, otherwise an exception is thrown
    try:
        sftpConnection.stat(pathToFile)
        return True
    except IOError:
        return False

# Parse the command line arguments we've been given
parser = argparse.ArgumentParser(description='Utility to transfer a given directory to a given location on a remote system over SFTP.')
parser.add_argument('--source', type=str, required=True)
parser.add_argument('--destination', type=str, required=True)
parser.add_argument('--server', type=str, required=True)
parser.add_argument('--username', type=str, required=True)
arguments = parser.parse_args()

# Let's establish a connection to the remote system
# We assume the private key file is password-less
privateKeyFile = open(os.environ['CRAFT_UPLOAD_KEY'], 'r')
uploadConnection = establishSSHConnection( arguments.server, arguments.username, privateKeyFile )
# Open the SFTP connection
sftp = uploadConnection.open_sftp()

# Step 1: Make sure the remote destination exists
# As a sanity mechanism, we don't continue if it does not in case there is a typo or something else (the remote admin should be creating it in advance for us)
if not sftpFileExists( sftp, arguments.destination ):
    print("Destination does not exist on the remote server - cannot continue!")
    sys.exit(1)

# Step 2: Start recursing through our local copy and uploading files
for root, directories, files in os.walk( arguments.source ):
    # First we need to make our root relative
    relativeRoot = os.path.relpath( root, arguments.source )
    # Now we can determine our remote prefix for this iteration
    remotePrefix = os.path.join( arguments.destination, relativeRoot )

    # First we check the directories to see if any need to be created
    for directory in directories:
        # Determine what it's full remote path would be
        # Because SFTP works with forward not backslashes we have to translate those before use
        remoteDirectory = os.path.join( remotePrefix, directory )
        remoteDirectory = remoteDirectory.replace('\\', '/')
        # Now check to see if it exists..
        if not sftpFileExists( sftp, remoteDirectory ):
            # And create it if it doesn't exist...
            sftp.mkdir( remoteDirectory )

    # Now we can handle files!
    for filename in files:
        # Determine what it's full local and remote paths are
        # Because SFTP works with forward not backslashes we have to translate those before use
        localPath = os.path.join( root, filename )
        remotePath = os.path.join( remotePrefix, filename )
        remotePath = remotePath.replace('\\', '/')

        # Prepare to begin checking whether we need to upload a file
        # By default we don't need to upload
        needToUpload = False

        # If any exception is raised we will assume we need to upload the file
        try:
            # First, let's stat the file on the server
            # If it doesn't exist this will fail with an exception, and we'll upload the file
            sftpStat = sftp.stat( remotePath )
            # We will need a local stat() to compare it to as well
            localStat = os.stat( localPath )
            # If the filesize differs, or if the local file is newer we need to upload
            if localStat.st_size != sftpStat.st_size or localStat.st_mtime > sftpStat.st_mtime:
                needToUpload = True

        # If anything went wrong then we need to upload!
        except IOError:
            needToUpload = True

        # Do we need to upload?
        if needToUpload:
            # Upload the file!
            sftp.put( localPath, remotePath )