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
parent
05e56b66a4
commit
39b877258b
|
@ -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,
|
||||
|
|
|
@ -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")
|
||||
|
|
Reference in New Issue