#!/usr/local/lib/python2.7
# Copyright (c) 2017 Arista Networks, Inc.  All rights reserved.
# Arista Networks, Inc. Confidential and Proprietary.

import paramiko
import argparse
import sys
import os

sftp_server = "<sftp_server>"
sftp_username = "<sftp_username>"
sftp_passphrase = "<sftp_passphrase>"

class SftpSession:
    """
    SFTP Session
    """

    def __init__(self, server, username, passphrase):
        self.sftp_port = 22
        self.key = paramiko.RSAKey.from_private_key_file(os.path.expanduser('~/.ssh/id_rsa'), password=self.passphrase)
        self._transport = paramiko.Transport(self.sftp_server, self.sftp_port)
        self._transport.connect(username=self.sftp_username, password=self.passphrase, pkey=self.key)
        self._sftp = paramiko.SFTPClient.from_transport(self._transport)

    def __enter__(self):
        return self

    def __exit__(self, *args):
        self.close()

    def close(self):
        self._transport.close()

    def get_dir(self, path):
        """Returns the list of files/directories at the path on the sftp server"""
        return self._sftp.listdir(path)

    def set_dir(self, path):
        """Creates a directory on the sftp server"""
        self._sftp.mkdir(path)

    def copy_to_sftp(self, local_path, sftp_path):
        """Copies the file to the specified path"""
        self._sftp.put(local_path, sftp_path)