khepri-tftpy/tftpy/TftpStates.py

from TftpShared import *
from TftpPacketTypes import *
from TftpPacketFactory import *
import socket, time

###############################################################################
# Utility classes
###############################################################################

class TftpMetrics(object):
    """A class representing metrics of the transfer."""
    def __init__(self):
        # Bytes transferred
        self.bytes = 0
        # Duplicate packets received
        self.dups = {}
        self.dupcount = 0
        # Times
        self.start_time = 0
        self.end_time = 0
        self.duration = 0
        # Rates
        self.bps = 0
        self.kbps = 0

    def compute(self):
        # Compute transfer time
        self.duration = self.end_time - self.start_time
        logger.debug("TftpMetrics.compute: duration is %s" % self.duration)
        self.bps = (self.bytes * 8.0) / self.duration
        self.kbps = self.bps / 1024.0
        logger.debug("TftpMetrics.compute: kbps is %s" % self.kbps)
        dupcount = 0
        for key in self.dups:
            dupcount += self.dups[key]

###############################################################################
# Context classes
###############################################################################

class TftpContext(object):
    """The base class of the contexts."""
    def __init__(self, host, port):
        """Constructor for the base context, setting shared instance
        variables."""
        self.factory = TftpPacketFactory()
        self.host = host
        self.port = port
        # The port associated with the TID
        self.tidport = None
        # Metrics
        self.metrics = TftpMetrics()

    def start(self):
        return NotImplementedError, "Abstract method"

    def end(self):
        return NotImplementedError, "Abstract method"

    def gethost(self):
        "Simple getter method for use in a property."
        return self.__host

    def sethost(self, host):
        """Setter method that also sets the address property as a result
        of the host that is set."""
        self.__host = host
        self.address = socket.gethostbyname(host)

    host = property(gethost, sethost)

    def sendAck(self, blocknumber):
        """This method sends an ack packet to the block number specified."""
        logger.info("sending ack to block %d" % blocknumber)
        ackpkt = TftpPacketACK()
        ackpkt.blocknumber = blocknumber
        self.sock.sendto(ackpkt.encode().buffer, (self.host, self.tidport))

    def sendError(self, errorcode):
        """This method uses the socket passed, and uses the errorcode to
        compose and send an error packet."""
        logger.debug("In sendError, being asked to send error %d" % errorcode)
        errpkt = TftpPacketERR()
        errpkt.errorcode = errorcode
        self.sock.sendto(errpkt.encode().buffer, (self.host, self.tidport))

class TftpContextClient(TftpContext):
    """This class represents shared functionality by both the download and
    upload client contexts."""
    def __init__(self, host, port, filename, options, packethook, timeout):
        TftpContext.__init__(self, host, port)
        self.file_to_transfer = filename
        self.options = options
        self.packethook = packethook
        self.sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        self.sock.settimeout(timeout)
        self.state = None
        self.next_block = 0

    def setNextBlock(self, block):
        if block > 2 ** 16:
            logger.debug("block number rollover to 0 again")
            block = 0
        self.__eblock = block

    def getNextBlock(self):
        return self.__eblock

    next_block = property(getNextBlock, setNextBlock)

    def cycle(self):
        """Here we wait for a response from the server after sending it
        something, and dispatch appropriate action to that response."""
        for i in range(TIMEOUT_RETRIES):
            logger.debug("in cycle, receive attempt %d" % i)
            try:
                (buffer, (raddress, rport)) = self.sock.recvfrom(MAX_BLKSIZE)
            except socket.timeout, err:
                logger.warn("Timeout waiting for traffic, retrying...")
                continue
            break
        else:
            raise TftpException, "Hit max timeouts, giving up."

        # Ok, we've received a packet. Log it.
        logger.debug("Received %d bytes from %s:%s"
                        % (len(buffer), raddress, rport))

        # Decode it.
        recvpkt = self.factory.parse(buffer)

        # Check for known "connection".
        if raddress != self.address:
            logger.warn("Received traffic from %s, expected host %s. Discarding"
                        % (raddress, self.host))

        if self.tidport and self.tidport != rport:
            logger.warn("Received traffic from %s:%s but we're "
                        "connected to %s:%s. Discarding."
                        % (raddress, rport,
                        self.host, self.tidport))

        # If there is a packethook defined, call it. We unconditionally
        # pass all packets, it's up to the client to screen out different
        # kinds of packets. This way, the client is privy to things like
        # negotiated options.
        if self.packethook:
            self.packethook(recvpkt)

        # And handle it, possibly changing state.
        self.state = self.state.handle(recvpkt, raddress, rport)

class TftpContextClientUpload(TftpContextClient):
    """The upload context for the client during an upload."""
    def __init__(self, host, port, filename, input, options, packethook, timeout):
        TftpContextClient.__init__(self,
                                   host,
                                   port,
                                   filename,
                                   options,
                                   packethook,
                                   timeout)
        self.fileobj = open(input, "wb")

        logger.debug("TftpContextClientUpload.__init__()")
        logger.debug("file_to_transfer = %s, options = %s" %
            (self.file_to_transfer, self.options))

    def start(self):
        logger.info("sending tftp upload request to %s" % self.host)
        logger.info("    filename -> %s" % self.file_to_transfer)
        logger.info("    options -> %s" % self.options)

        self.metrics.start_time = time.time()
        logger.debug("set metrics.start_time to %s" % self.metrics.start_time)

        # FIXME: put this in a sendWRQ method?
        pkt = TftpPacketWRQ()
        pkt.filename = self.file_to_transfer
        pkt.mode = "octet" # FIXME - shouldn't hardcode this
        pkt.options = self.options
        self.sock.sendto(pkt.encode().buffer, (self.host, self.port))
        self.next_block = 1

        self.state = TftpStateSentWRQ(self)

        try:
            while self.state:
                logger.debug("state is %s" % self.state)
                self.cycle()
        finally:
            self.fileobj.close()

    def end(self):
        pass

class TftpContextClientDownload(TftpContextClient):
    """The download context for the client during a download."""
    def __init__(self, host, port, filename, output, options, packethook, timeout):
        TftpContextClient.__init__(self,
                                   host,
                                   port,
                                   filename,
                                   options,
                                   packethook,
                                   timeout)
        # FIXME - need to support alternate return formats than files?
        # File-like objects would be ideal, ala duck-typing.
        self.fileobj = open(output, "wb")

        logger.debug("TftpContextClientDownload.__init__()")
        logger.debug("file_to_transfer = %s, options = %s" %
            (self.file_to_transfer, self.options))

    def start(self):
        """Initiate the download."""
        logger.info("sending tftp download request to %s" % self.host)
        logger.info("    filename -> %s" % self.file_to_transfer)
        logger.info("    options -> %s" % self.options)

        self.metrics.start_time = time.time()
        logger.debug("set metrics.start_time to %s" % self.metrics.start_time)

        # FIXME: put this in a sendRRQ method?
        pkt = TftpPacketRRQ()
        pkt.filename = self.file_to_transfer
        pkt.mode = "octet" # FIXME - shouldn't hardcode this
        pkt.options = self.options
        self.sock.sendto(pkt.encode().buffer, (self.host, self.port))
        self.next_block = 1

        self.state = TftpStateSentRRQ(self)

        try:
            while self.state:
                logger.debug("state is %s" % self.state)
                self.cycle()
        finally:
            self.fileobj.close()

    def end(self):
        """Finish up the context."""
        self.metrics.end_time = time.time()
        logger.debug("set metrics.end_time to %s" % self.metrics.end_time)
        self.metrics.compute()


###############################################################################
# State classes
###############################################################################

class TftpState(object):
    """The base class for the states."""

    def __init__(self, context):
        """Constructor for setting up common instance variables. The involved
        file object is required, since in tftp there's always a file
        involved."""
        self.context = context

    def handle(self, pkt, raddress, rport):
        """An abstract method for handling a packet. It is expected to return
        a TftpState object, either itself or a new state."""
        raise NotImplementedError, "Abstract method"

    def handleOACK(self, pkt):
        """This method handles an OACK from the server, syncing any accepted
        options."""
        if pkt.options.keys() > 0:
            if pkt.match_options(self.context.options):
                logger.info("Successful negotiation of options")
                # Set options to OACK options
                self.context.options = pkt.options
                for key in self.context.options:
                    logger.info("    %s = %s" % (key, self.context.options[key]))
            else:
                logger.error("failed to negotiate options")
                raise TftpException, "Failed to negotiate options"
        else:
            raise TftpException, "No options found in OACK"

class TftpStateUpload(TftpState):
    """A class holding common code for upload states."""
    def sendDat(self, resend=False):
        finished = False
        blocknumber = self.context.next_block
        if not resend:
            blksize = int(self.context.options['blksize'])
            buffer = self.context.fileobj.read(blksize)
            logger.debug("Read %d bytes into buffer" % len(buffer))
            if len(buffer) < blksize:
                logger.info("Reached EOF on file %s" % self.context.input)
                finished = True
            self.context.next_block += 1
            self.bytes += len(buffer)
        else:
            logger.warn("Resending block number %d" % blocknumber)
        dat = TftpPacketDAT()
        dat.data = buffer
        dat.blocknumber = blocknumber
        logger.debug("Sending DAT packet %d" % blocknumber)
        self.context.sock.sendto(dat.encode().buffer,
                                 (self.context.host, self.context.port))
        if self.context.packethook:
            self.context.packethook(dat)
        return finished

class TftpStateDownload(TftpState):
    """A class holding common code for download states."""
    def handleDat(self, pkt):
        """This method handles a DAT packet during a download."""
        logger.info("handling DAT packet - block %d" % pkt.blocknumber)
        logger.debug("expecting block %s" % self.context.next_block)
        if pkt.blocknumber == self.context.next_block:
            logger.debug("good, received block %d in sequence"
                        % pkt.blocknumber)

            self.context.sendAck(pkt.blocknumber)
            self.context.next_block += 1

            logger.debug("writing %d bytes to output file"
                        % len(pkt.data))
            self.context.fileobj.write(pkt.data)
            self.context.metrics.bytes += len(pkt.data)
            # Check for end-of-file, any less than full data packet.
            if len(pkt.data) < int(self.context.options['blksize']):
                logger.info("end of file detected")
                return None

        elif pkt.blocknumber < self.context.next_block:
            logger.warn("dropping duplicate block %d" % pkt.blocknumber)
            if self.context.metrics.dups.has_key(pkt.blocknumber):
                self.context.metrics.dups[pkt.blocknumber] += 1
            else:
                self.context.metrics.dups[pkt.blocknumber] = 1
            tftpassert(self.context.metrics.dups[pkt.blocknumber] < MAX_DUPS,
                    "Max duplicates for block %d reached" % pkt.blocknumber)
            # FIXME: double-check sorceror's apprentice problem!
            logger.debug("ACKing block %d again, just in case" % pkt.blocknumber)
            self.context.sendAck(pkt.blocknumber)

        else:
            # FIXME: should we be more tolerant and just discard instead?
            msg = "Whoa! Received future block %d but expected %d" \
                % (pkt.blocknumber, self.context.next_block)
            logger.error(msg)
            raise TftpException, msg

        # Default is to ack
        return TftpStateSentACK(self.context)

class TftpStateSentWRQ(TftpStateUpload):
    """Just sent an WRQ packet for an upload."""
    def handle(self, pkt, raddress, rport):
        """Handle a packet we just received."""
        if not self.context.tidport:
            self.context.tidport = rport
            logger.debug("Set remote port for session to %s" % rport)

        # If we're going to successfully transfer the file, then we should see
        # either an OACK for accepted options, or an ACK to ignore options.
        if isinstance(pkt, TftpPacketOACK):
            logger.info("received OACK from server")
            try:
                self.handleOACK(pkt)
            except TftpException, err:
                logger.error("failed to negotiate options")
                self.context.sendError(TftpErrors.FailedNegotiation)
                raise
            else:
                logger.debug("sending first DAT packet")
                fin = self.context.sendDat()
                if fin:
                    logger.info("Add done")
                    return None
                else:
                    logger.debug("Changing state to TftpStateSentDAT")
                    return TftpStateSentDAT(self.context)

        elif isinstance(pkt, TftpPacketACK):
            logger.info("received ACK from server")
            logger.debug("apparently the server ignored our options")
            # The block number should be zero.
            if pkt.blocknumber == 0:
                logger.debug("ack blocknumber is zero as expected")
                logger.debug("sending first DAT packet")
                fin = self.context.sendDat()
                if fin:
                    logger.info("Add done")
                    return None
                else:
                    logger.debug("Changing state to TftpStateSentDAT")
                    return TftpStateSentDAT(self.context)
            else:
                logger.warn("discarding ACK to block %s" % pkt.blocknumber)
                logger.debug("still waiting for valid response from server")
                return self

        elif isinstance(pkt, TftpPacketERR):
            self.context.sendError(TftpErrors.IllegalTftpOp)
            raise TftpException, "Received ERR from server: " + str(pkt)

        elif isinstance(pkt, TftpPacketRRQ):
            self.context.sendError(TftpErrors.IllegalTftpOp)
            raise TftpException, "Received RRQ from server while in upload"

        elif isinstance(pkt, TftpPacketDAT):
            self.context.sendError(TftpErrors.IllegalTftpOp)
            raise TftpException, "Received DAT from server while in upload"

        else:
            self.context.sendError(TftpErrors.IllegalTftpOp)
            raise TftpException, "Received unknown packet type from server: " + str(pkt)

        # By default, no state change.
        return self

class TftpStateSentDAT(TftpStateUpload):
    """This class represents the state of the transfer when a DAT was just
    sent, and we are waiting for an ACK from the server. This class is the
    same one used by the client during the upload, and the server during the
    download."""

class TftpStateSentRRQ(TftpStateDownload):
    """Just sent an RRQ packet."""
    def handle(self, pkt, raddress, rport):
        """Handle the packet in response to an RRQ to the server."""
        if not self.context.tidport:
            self.context.tidport = rport
            logger.debug("Set remote port for session to %s" % rport)

        # Now check the packet type and dispatch it properly.
        if isinstance(pkt, TftpPacketOACK):
            logger.info("received OACK from server")
            try:
                self.handleOACK(pkt)
            except TftpException, err:
                logger.error("failed to negotiate options: %s" % str(err))
                self.context.sendError(TftpErrors.FailedNegotiation)
                raise
            else:
                logger.debug("sending ACK to OACK")

                self.context.sendAck(blocknumber=0)

                logger.debug("Changing state to TftpStateSentACK")
                return TftpStateSentACK(self.context)

        elif isinstance(pkt, TftpPacketDAT):
            # If there are any options set, then the server didn't honour any
            # of them.
            logger.info("received DAT from server")
            if self.context.options:
                logger.info("server ignored options, falling back to defaults")
                self.context.options = { 'blksize': DEF_BLKSIZE }
            return self.handleDat(pkt)

        # Every other packet type is a problem.
        elif isinstance(recvpkt, TftpPacketACK):
            # Umm, we ACK, the server doesn't.
            self.context.sendError(TftpErrors.IllegalTftpOp)
            raise TftpException, "Received ACK from server while in download"

        elif isinstance(recvpkt, TftpPacketWRQ):
            self.context.sendError(TftpErrors.IllegalTftpOp)
            raise TftpException, "Received WRQ from server while in download"

        elif isinstance(recvpkt, TftpPacketERR):
            self.context.sendError(TftpErrors.IllegalTftpOp)
            raise TftpException, "Received ERR from server: " + str(recvpkt)

        else:
            self.context.sendError(TftpErrors.IllegalTftpOp)
            raise TftpException, "Received unknown packet type from server: " + str(recvpkt)

        # By default, no state change.
        return self

class TftpStateSentACK(TftpStateDownload):
    """Just sent an ACK packet. Waiting for DAT."""
    def handle(self, pkt, raddress, rport):
        """Handle the packet in response to an ACK, which should be a DAT."""
        if isinstance(pkt, TftpPacketDAT):
            return self.handleDat(pkt)

        # Every other packet type is a problem.
        elif isinstance(recvpkt, TftpPacketACK):
            # Umm, we ACK, the server doesn't.
            self.context.sendError(TftpErrors.IllegalTftpOp)
            raise TftpException, "Received ACK from server while in download"

        elif isinstance(recvpkt, TftpPacketWRQ):
            self.context.sendError(TftpErrors.IllegalTftpOp)
            raise TftpException, "Received WRQ from server while in download"

        elif isinstance(recvpkt, TftpPacketERR):
            self.context.sendError(TftpErrors.IllegalTftpOp)
            raise TftpException, "Received ERR from server: " + str(recvpkt)

        else:
            self.context.sendError(TftpErrors.IllegalTftpOp)
            raise TftpException, "Received unknown packet type from server: " + str(recvpkt)