Adding support for input/output as stdin/stdout

This commit is contained in:
Michael P. Soulier 2010-05-25 21:51:07 -04:00
parent f4a3ff6356
commit 58623df7d0
3 changed files with 20 additions and 8 deletions

View file

@ -29,10 +29,10 @@ def main():
default=512) default=512)
parser.add_option('-o', parser.add_option('-o',
'--output', '--output',
help='output file (default: same as requested filename)') help='output file, - for stdout (default: same as download)')
parser.add_option('-i', parser.add_option('-i',
'--input', '--input',
help='input file (default: same as upload filename)') help='input file, - for stdin (default: same as upload)')
parser.add_option('-d', parser.add_option('-d',
'--debug', '--debug',
action='store_true', action='store_true',

View file

@ -33,7 +33,9 @@ class TftpClient(TftpSession):
copy of each DAT packet received in the form of a TftpPacketDAT copy of each DAT packet received in the form of a TftpPacketDAT
object. The timeout parameter may be used to override the default object. The timeout parameter may be used to override the default
SOCK_TIMEOUT setting, which is the amount of time that the client will SOCK_TIMEOUT setting, which is the amount of time that the client will
wait for a receive packet to arrive.""" wait for a receive packet to arrive.
Note: If output is a hyphen then stdout is used."""
# We're downloading. # We're downloading.
log.debug("Creating download context with the following params:") log.debug("Creating download context with the following params:")
log.debug("host = %s, port = %s, filename = %s, output = %s" log.debug("host = %s, port = %s, filename = %s, output = %s"
@ -64,6 +66,7 @@ class TftpClient(TftpSession):
log.info("Received %d duplicate packets" % metrics.dupcount) log.info("Received %d duplicate packets" % metrics.dupcount)
def upload(self, filename, input, packethook=None, timeout=SOCK_TIMEOUT): def upload(self, filename, input, packethook=None, timeout=SOCK_TIMEOUT):
"""Note: If input is a hyphen then stdin is used."""
# Open the input file. # Open the input file.
# FIXME: As of the state machine, this is now broken. Need to # FIXME: As of the state machine, this is now broken. Need to
# implement with new state machine. # implement with new state machine.

View file

@ -1,7 +1,7 @@
from TftpShared import * from TftpShared import *
from TftpPacketTypes import * from TftpPacketTypes import *
from TftpPacketFactory import * from TftpPacketFactory import *
import socket, time, os import socket, time, os, sys
############################################################################### ###############################################################################
# Utility classes # Utility classes
@ -223,7 +223,8 @@ class TftpContextServer(TftpContext):
self.metrics.compute() self.metrics.compute()
class TftpContextClientUpload(TftpContext): class TftpContextClientUpload(TftpContext):
"""The upload context for the client during an upload.""" """The upload context for the client during an upload.
Note: If input is a hyphen, then we will use stdin."""
def __init__(self, def __init__(self,
host, host,
port, port,
@ -239,7 +240,10 @@ class TftpContextClientUpload(TftpContext):
self.file_to_transfer = filename self.file_to_transfer = filename
self.options = options self.options = options
self.packethook = packethook self.packethook = packethook
self.fileobj = open(input, "rb") if input == '-':
self.fileobj = sys.stdin
else:
self.fileobj = open(input, "rb")
log.debug("TftpContextClientUpload.__init__()") log.debug("TftpContextClientUpload.__init__()")
log.debug("file_to_transfer = %s, options = %s" % log.debug("file_to_transfer = %s, options = %s" %
@ -278,7 +282,8 @@ class TftpContextClientUpload(TftpContext):
self.metrics.compute() self.metrics.compute()
class TftpContextClientDownload(TftpContext): class TftpContextClientDownload(TftpContext):
"""The download context for the client during a download.""" """The download context for the client during a download.
Note: If output is a hyphen, then the output will be sent to stdout."""
def __init__(self, def __init__(self,
host, host,
port, port,
@ -297,7 +302,11 @@ class TftpContextClientDownload(TftpContext):
self.packethook = packethook self.packethook = packethook
# FIXME - need to support alternate return formats than files? # FIXME - need to support alternate return formats than files?
# File-like objects would be ideal, ala duck-typing. # File-like objects would be ideal, ala duck-typing.
self.fileobj = open(output, "wb") # If the filename is -, then use stdout
if output == '-':
self.fileobj = sys.stdout
else:
self.fileobj = open(output, "wb")
log.debug("TftpContextClientDownload.__init__()") log.debug("TftpContextClientDownload.__init__()")
log.debug("file_to_transfer = %s, options = %s" % log.debug("file_to_transfer = %s, options = %s" %