First working upload with new state machine. Not usable yet, upload fails to

always send all data for some reason.

3 files changed, 40 insertions, 12 deletions

diff --git a/bin/tftpy_client.py b/bin/tftpy_client.py
index ac09397..039e7d3 100755
--- a/bin/tftpy_client.py
+++ b/bin/tftpy_client.py
@@ -16,7 +16,10 @@ def main():
                       default=69)
     parser.add_option('-f',
                       '--filename',
-                      help='filename to fetch')
+                      help='filename to fetch (deprecated, use download)')
+    parser.add_option('-D',
+                      '--download',
+                      help='filename to download')
     parser.add_option('-u',
                       '--upload',
                       help='filename to upload')
@@ -46,7 +49,10 @@ def main():
                       default=False,
                       help="ask client to send tsize option in download")
     options, args = parser.parse_args()
-    if not options.host or (not options.filename and not options.upload):
+    # Handle legacy --filename argument.
+    if options.filename:
+        options.download = options.filename
+    if not options.host or (not options.download and not options.upload):
         sys.stderr.write("Both the --host and --filename options "
                          "are required.\n")
         parser.print_help()
@@ -88,10 +94,10 @@ def main():
                                int(options.port),
                                tftp_options)
     try:
-        if options.filename:
+        if options.download:
             if not options.output:
-                options.output = os.path.basename(options.filename)
-            tclient.download(options.filename,
+                options.output = os.path.basename(options.download)
+            tclient.download(options.download,
                             options.output,
                             progresshook)
         elif options.upload:
diff --git a/tftpy/TftpServer.py b/tftpy/TftpServer.py
index fe8a96e..862f1a1 100644
--- a/tftpy/TftpServer.py
+++ b/tftpy/TftpServer.py
@@ -56,7 +56,7 @@ class TftpServer(TftpSession):
         log.info("Server requested on ip %s, port %s"
                 % (listenip, listenport))
         try:
-            # FIXME - sockets should be non-blocking?
+            # FIXME - sockets should be non-blocking
             self.sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
             self.sock.bind((listenip, listenport))
         except socket.error, err:
@@ -106,6 +106,9 @@ class TftpServer(TftpSession):
                     else:
                         log.warn("received traffic on main socket for "
                                  "existing session??")
+                    log.info("Currently handling these sessions:")
+                    for session_key, session in self.sessions.items():
+                        log.info("    %s" % session)
 
                 else:
                     # Must find the owner of this traffic.
@@ -156,6 +159,7 @@ class TftpServer(TftpSession):
                     log.info("%d duplicate packets" % metrics.dupcount)
                     log.debug("Deleting session %s" % key)
                     del self.sessions[key]
+                    log.debug("Session list is now %s" % self.sessions)
                 else:
                     log.warn("Strange, session %s is not on the deletion list"
                         % key)
diff --git a/tftpy/TftpStates.py b/tftpy/TftpStates.py
index 942b073..e26b268 100644
--- a/tftpy/TftpStates.py
+++ b/tftpy/TftpStates.py
@@ -179,7 +179,10 @@ class TftpContextServer(TftpContext):
         self.dyn_file_func = dyn_file_func
         # In a server, the tidport is the same as the port. This is also true
         # with symmetric UDP, which we haven't implemented yet.
-        self.tidport = port
+        #self.tidport = port
+
+    def __str__(self):
+        return "%s:%s %s" % (self.host, self.port, self.state)
 
     def start(self, buffer):
         """Start the state cycle. Note that the server context receives an
@@ -227,12 +230,15 @@ class TftpContextClientUpload(TftpContext):
         self.file_to_transfer = filename
         self.options = options
         self.packethook = packethook
-        self.fileobj = open(input, "wb")
+        self.fileobj = open(input, "rb")
 
         log.debug("TftpContextClientUpload.__init__()")
         log.debug("file_to_transfer = %s, options = %s" %
             (self.file_to_transfer, self.options))
 
+    def __str__(self):
+        return "%s:%s %s" % (self.host, self.port, self.state)
+
     def start(self):
         log.info("Sending tftp upload request to %s" % self.host)
         log.info("    filename -> %s" % self.file_to_transfer)
@@ -290,6 +296,9 @@ class TftpContextClientDownload(TftpContext):
         log.debug("file_to_transfer = %s, options = %s" %
             (self.file_to_transfer, self.options))
 
+    def __str__(self):
+        return "%s:%s %s" % (self.host, self.port, self.state)
+
     def start(self):
         """Initiate the download."""
         log.info("Sending tftp download request to %s" % self.host)
@@ -392,6 +401,9 @@ class TftpState(object):
         it is required to send an OACK to the client."""
         options = pkt.options
         sendoack = False
+        if not self.context.tidport:
+            self.context.tidport = rport
+            log.info("Setting tidport to %s" % rport)
         if not options:
             log.debug("Setting default options, blksize")
             # FIXME: put default options elsewhere
@@ -458,7 +470,7 @@ class TftpState(object):
         self.context.metrics.bytes += len(dat.data)
         log.debug("Sending DAT packet %d" % dat.blocknumber)
         self.context.sock.sendto(dat.encode().buffer,
-                                 (self.context.host, self.context.port))
+                                 (self.context.host, self.context.tidport))
         if self.context.packethook:
             self.context.packethook(dat)
         self.context.last_dat_pkt = dat
@@ -566,6 +578,7 @@ class TftpStateServerRecvRRQ(TftpState):
             # acknowledgement to an OACK.
             # FIXME: perhaps we do need a TftpStateExpectOACK class...
             self.sendOACK()
+            # Note, self.context.next_block is already 0.
         else:
             self.context.next_block = 1
             log.debug("No requested options, starting send...")
@@ -598,8 +611,11 @@ class TftpStateServerRecvWRQ(TftpState):
             log.debug("Sending OACK to client")
             self.sendOACK()
         else:
-            log.debug("No requested options, starting send...")
+            log.debug("No requested options, expecting transfer to begin...")
             self.sendACK()
+        # Whether we're sending an oack or not, we're expecting a DAT for
+        # block 1
+        self.context.next_block = 1
         # We may have sent an OACK, but we're expecting a DAT as the response
         # to either the OACK or an ACK, so lets unconditionally use the
         # TftpStateExpectDAT state.
@@ -609,7 +625,9 @@ class TftpStateServerRecvWRQ(TftpState):
         # up to the caller.
 
 class TftpStateServerStart(TftpState):
-    """The start state for the server."""
+    """The start state for the server. This is a transitory state since at
+    this point we don't know if we're handling an upload or a download. We
+    will commit to one of them once we interpret the initial packet."""
     def handle(self, pkt, raddress, rport):
         """Handle a packet we just received."""
         log.debug("In TftpStateServerStart.handle")
@@ -754,7 +772,7 @@ class TftpStateSentRRQ(TftpState):
         """Handle the packet in response to an RRQ to the server."""
         if not self.context.tidport:
             self.context.tidport = rport
-            log.debug("Set remote port for session to %s" % rport)
+            log.info("Set remote port for session to %s" % rport)
 
         # Now check the packet type and dispatch it properly.
         if isinstance(pkt, TftpPacketOACK):