Pure Python SCP library based on Twisted SSH libraries.

scp.py

#!/usr/bin/env python

"""
scp.py - SCP implementation using Twisted.

This currently only uploads files to the remote host. Receiving files and
uploading directories is NYI.

Credit: http://bit.ly/1EG3eL4

>>> import scp
>>> scp.upload(hostname='example.com', username='jathan', password='password',
               src_path='/etc/hosts', dst_path='/tmp/hosts')
True

"""

__author__ = 'Jathan McCollum'
__maintainer__ = 'Jathan McCollum'
__email__ = 'jathan@dropbox.com'
__copyright__ = 'Copyright (c) 2015 Dropbox, Inc.'
__version__ = '1.0'
__license__ = 'APL-2.0'


import getpass
import os
import sys
from twisted.conch.client.knownhosts import KnownHostsFile
from twisted.conch.endpoints import SSHCommandClientEndpoint
from twisted.conch.ssh.keys import EncryptedKeyError, Key
from twisted.internet import defer, protocol
from twisted.internet.endpoints import UNIXClientEndpoint
from twisted.internet.task import react
from twisted.python import log
from twisted.python.filepath import FilePath


# Dat debug doe
if os.getenv('DEBUG'):
    log.startLogging(sys.stderr)


class ScpProtocol(protocol.Protocol, object):
    """SCP Protocol. Currently only supports uploading files."""
    state = None  # Used by internal state machine for data transfer
    todo = 0  # Used to count remaining bytes for data transfer
    buf = ''  # Used when receiving files (NYI)
    chunk_size = 2 ** 14  # How many bytes to read at a time

    #: Default permissions. This is required when uploading via SCP.
    permissions = oct(0644)  # -rw-r--r--

    def __init__(self, src_path, dst_path, permissions=None):
        self.src_path = os.path.expanduser(src_path)
        self.src_size = os.path.getsize(self.src_path)
        self.src_data = open(self.src_path, 'rb')
        self.dst_path = dst_path

        # Split dst_path into (dir, filename)
        dst_dir, dst_file = os.path.split(dst_path)
        self.dst_dir = dst_dir
        self.dst_file = dst_file

        # self.command = 'scp -f %s' % (dst_dir,)  # Copy FROM remote host
        self.command = 'scp -t %s' % (dst_dir,)  # Copy TO remote host
        log.msg('COMMAND = %r *****' % (self.command,))

        # Octal permissions
        if permissions is None:
            permissions = self.permissions

        super(ScpProtocol, self).__init__()

    def __call__(self):
        """This is so we can play nicely w/ Factory objects."""
        return self

    def connectionMade(self):
        log.msg('CONNECTION MADE *****')
        self.finished = defer.Deferred()
        self.sendSource()

    def sendSource(self):
        """
        Send the initial source negotiation.

        See: https://blogs.oracle.com/janp/entry/how_the_scp_protocol_works
        """
        # Format: "Coctalperms size filename\n"
        msg = 'C%s %s %s\n' % (self.permissions, self.src_size, self.dst_file)
        log.msg('SENDING %r *****' % msg)
        self.transport.write(msg)
        self.state = 'sending'

    def connectionLost(self, reason):
        log.msg('CONNECTION LOST *****')
        self.src_data.close()  # Close the upload file handle
        self.finished.callback(None)

    def get_chunk(self):
        """Read a chunk from the source data."""
        return self.src_data.read(self.chunk_size)

    def dataReceived(self, data):
        log.msg('dataReceived: %s %r' % (self.state, data))

        # Uploading a file to remote side.
        if self.state == 'sending':
            self.todo = self.src_size
            log.msg('TODO: %r' % self.todo)

            while self.todo > 0:
                chunk = self.get_chunk()
                self.transport.write(chunk)
                self.todo -= self.chunk_size
            else:
                log.msg('UPLOAD DONE *****')
                self.transport.loseConnection()

        else:
            log.err("dataReceived in unknown state: %r" % (self.state,))

        # NYI
        '''
        # Waiting for remote file transfer to start
        elif self.state == 'waiting':
            # we've started the transfer, and are expecting a C
            # Coctalperms size filename\n

            # might not get it all at once, buffer
            self.buf += data
            if not self.buf.endswith('\n'):
                return

            b = self.buf
            self.buf = ''

            # Must be a C
            if not b.startswith('C'):
                log.msg("expecting C command: %r" % (self.buf,))
                self.loseConnection()
                return

            # Get the file info
            p, l, n = b[1:-1].split(' ')
            perms = int(p, 8)
            self.todo = int(l)
            log.msg("getting file %s mode %s len %i" % (n, oct(perms),
                    self.todo))

            # Tell the far end to start sending the content
            self.state = 'receiving'
            self.transport.write('\0')

        # Receiving a file from remote system.
        elif self.state == 'receiving':
            #log.msg('got %i bytes' % (len(data),))

            if len(data) > self.todo:
                extra = data[self.todo:]
                data = data[:self.todo]

                if extra != '\0':
                    log.msg("got %i more bytes than we expected, ignoring: %r"
                            % (len(extra), extra))

            DST.write(data)
            self.todo -= len(data)

            if self.todo <= 0:
                log.msg('done')
                self.loseConnection()
        '''


def read_key(path):
    """Read an SSH private key file."""
    try:
        return Key.fromFile(path)
    except EncryptedKeyError:
        passphrase = getpass.getpass("%r keyphrase: " % (path,))
        return Key.fromFile(path, passphrase=passphrase)


def perform_upload(reactor, hostname, username, password, src_path, dst_path,
                   port=22, agent=None, known_hosts=None, identity=None):
    """Upload a file to remote host using SCP."""

    # Keys
    keys = []
    if identity is not None:
        key_path = os.path.expanduser(identity)
        if os.path.exists(key_path):
            keys.append(read_key(key_path))

    # Setup known_hosts
    if known_hosts is None:
        known_hosts = '~/.ssh/known_hosts'
    known_hosts_path = FilePath(os.path.expanduser(known_hosts))
    if known_hosts_path.exists():
        knownHosts = KnownHostsFile.fromPath(known_hosts_path)
    else:
        knownHosts = None

    # Setup ssh-agent
    if agent is None or 'SSH_AUTH_SOCK' not in os.environ:
        agentEndpoint = None
    else:
        agentEndpoint = UNIXClientEndpoint(
            reactor, os.environ['SSH_AUTH_SOCK']
        )

    # Setup protocol and factory
    scp_protocol = ScpProtocol(src_path=src_path, dst_path=dst_path)
    factory = protocol.Factory.forProtocol(scp_protocol)

    # Setup endpoint
    endpoint = SSHCommandClientEndpoint.newConnection(
        reactor=reactor,
        command=scp_protocol.command,  # e.g. 'scp -t filename'
        username=username,
        hostname=hostname,
        port=port,
        password=password,
        knownHosts=knownHosts,
        keys=keys,
        agentEndpoint=agentEndpoint,
    )

    d = endpoint.connect(factory)
    d.addCallback(lambda proto: proto.finished)

    return d


def upload(hostname, username, password, src_path, dst_path,
           port=22, agent=None, known_hosts=None, identity=None):
    """This is so you can write a script to do stuff."""
    argv = [
        hostname, username, password, src_path, dst_path, port, agent,
        known_hosts, identity
    ]

    try:
        react(perform_upload, argv)
    except SystemExit as err:
        return err.code == 0


def usage():
    sys.exit("%s: src_path [user[:pass]@]hostname:dst_path" % (sys.argv[0],))


def main():
    args = sys.argv[1:]
    if len(args) < 2:
        usage()

    src_path = args[0]
    dst_info = args[1]

    # Parse username, password, hostname, and dst_path from dst_info
    if '@' in dst_info:
        username, dst_info = dst_info.split('@', 1)
    else:
        username = getpass.getuser()

    if ':' not in dst_info:
        usage()

    hostname, dst_path = dst_info.split(':', 1)

    if ':' in username:
        username, password = username.split(':', 1)
    else:
        password = getpass.getpass(
            'password for %s@%s: ' % (username, hostname)
        )

    argv = [hostname, username, password, src_path, dst_path]
    react(perform_upload, argv)


if __name__ == '__main__':
    main()