Created
August 21, 2015 00:02
-
-
Save jathanism/1cf8f1d73c6711dfcc2f to your computer and use it in GitHub Desktop.
Pure Python SCP library based on Twisted SSH libraries.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
#!/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() |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment