Allow file objects to be used for download output or upload input.

Objects are considered file-like if they have read() or write()
functions. If they do not, they are assumed to be strings to a local
path (existing behavior).
master
Jay Weisskopf 2012-05-03 01:24:25 -05:00
parent 05e56b66a4
commit 39b877258b
2 changed files with 27 additions and 23 deletions

View File

@ -27,15 +27,15 @@ class TftpClient(TftpSession):
def download(self, filename, output, packethook=None, timeout=SOCK_TIMEOUT):
"""This method initiates a tftp download from the configured remote
host, requesting the filename passed. It saves the file to a local
file specified in the output parameter. If a packethook is provided,
it must be a function that takes a single parameter, which will be a
copy of each DAT packet received in the form of a TftpPacketDAT
object. The timeout parameter may be used to override the default
SOCK_TIMEOUT setting, which is the amount of time that the client will
wait for a receive packet to arrive.
host, requesting the filename passed. It writes the file to output,
which can be a file-like object or a path to a local file. If a
packethook is provided, it must be a function that takes a single
parameter, which will be a copy of each DAT packet received in the
form of a TftpPacketDAT object. The timeout parameter may be used to
override the default SOCK_TIMEOUT setting, which is the amount of time
that the client will wait for a receive packet to arrive.
Note: If output is a hyphen then stdout is used."""
Note: If output is a hyphen, stdout is used."""
# We're downloading.
log.debug("Creating download context with the following params:")
log.debug("host = %s, port = %s, filename = %s, output = %s"
@ -67,17 +67,15 @@ class TftpClient(TftpSession):
def upload(self, filename, input, packethook=None, timeout=SOCK_TIMEOUT):
"""This method initiates a tftp upload to the configured remote host,
uploading the filename passed. If a packethook is provided, it must
be a function that takes a single parameter, which will be a copy of
each DAT packet sent in the form of a TftpPacketDAT object. The
timeout parameter may be used to override the default SOCK_TIMEOUT
setting, which is the amount of time that the client will wait for a
DAT packet to be ACKd by the server.
uploading the filename passed. It reads the file from input, which
can be a file-like object or a path to a local file. If a packethook
is provided, it must be a function that takes a single parameter,
which will be a copy of each DAT packet sent in the form of a
TftpPacketDAT object. The timeout parameter may be used to override
the default SOCK_TIMEOUT setting, which is the amount of time that
the client will wait for a DAT packet to be ACKd by the server.
The input option is the full path to the file to upload, which can
optionally be '-' to read from stdin.
Note: If output is a hyphen then stdout is used."""
Note: If input is a hyphen, stdin is used."""
self.context = TftpContextClientUpload(self.host,
self.iport,
filename,

View File

@ -255,7 +255,11 @@ class TftpContextClientUpload(TftpContext):
self.file_to_transfer = filename
self.options = options
self.packethook = packethook
if input == '-':
# If the input object has a read() function,
# assume it is file-like.
if hasattr(input, 'read'):
self.fileobj = input
elif input == '-':
self.fileobj = sys.stdin
else:
self.fileobj = open(input, "rb")
@ -328,10 +332,12 @@ class TftpContextClientDownload(TftpContext):
self.file_to_transfer = filename
self.options = options
self.packethook = packethook
# FIXME - need to support alternate return formats than files?
# File-like objects would be ideal, ala duck-typing.
# If the filename is -, then use stdout
if output == '-':
# If the output object has a write() function,
# assume it is file-like.
if hasattr(output, 'write'):
self.fileobj = output
# If the output filename is -, then use stdout
elif output == '-':
self.fileobj = sys.stdout
else:
self.fileobj = open(output, "wb")