#!/usr/bin/python
# -*- Mode: Python; py-indent-offset: 4 -*-
#
# Copyright (C) 2001,2007  Ray Burr
# 
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
# 
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
# 
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
# THE SOFTWARE.
# 

'''
A TFTP client.

:authors: Ray Burr
:license: MIT License
:contact: http://www.nightmare.com/~ryb/

Simple use:

    # Copy the file remote.txt on somehost to the local file local.txt.
    f = open("local.txt", "w")
    t = TftpDownloader(f, "somehost", "remote.txt")
    t.transfer()

    # An even easier way for the same.
    get("local.txt", "somehost", "remote.txt")

or

    # Copy the local file local.txt to the file remote.txt on somehost.
    f = open("local.txt", "r")
    t = TftpUploader(f, "somehost", "remote.txt")
    t.transfer()

I attempted to copy the logic used in the BSD tftp client.  That
client (and others based on it) does not follow the standards
(RFC-1350 and RFC-1123) in a few ways.  It does not fix the
"Sorcerer's Apprentice Syndrome" bug, and has a simple retransmission
timeout algorithm.  It should be reliable for use on a LAN but
probably not for the Internet.

This version only supports the "OCTET" transfer mode.  "NETASCII" is
not supported.
''' #"

__version__ = "20070611"

from struct import pack, unpack
import fcntl
import select
import socket
import string
import sys
import time
import termios

RRQ   = 1
WRQ   = 2
DATA  = 3
ACK   = 4
ERROR = 5

MAX_PACKET_SIZE = 512 + 4


class Error(Exception):
    pass


class ProtocolError(Error):
    pass


class TftpBase:

    def __init__(self, name, **options):
        self.name = name
        self.mode = options.get("mode")
        if self.mode is None:
            self.mode = "octet"
        if string.lower(self.mode) != "octet":
            raise Error(
                "The specified TFTP transfer mode %s"
                " is not implemented." % repr(self.mode))
        self.tftpInitialPort = options.get("port")
        if self.tftpInitialPort is None:
            self.tftpInitialPort = socket.getservbyname("tftp", "udp")
        self.resendTimeout = options.get("resendTimeout")
        if self.resendTimeout is None:
            self.resendTimeout = 5.0
        self.maxTimeout = options.get("maxTimeout")
        if self.maxTimeout is None:
            self.maxTimeout = self.resendTimeout * 5
        self.trace = options.get("trace", 0)
        self.segmentSize = 512
        self.state = "start"
        self.done = 0
        self.serverPort = None
        self.needAck = 0
        self.needResend = 0
        self.lastSegment = 0
        self.blockNumber = 0
        self.lastPacket = None
        self.lastPacketTime = None
        self.lastResendTime = None

    def makeWriteRequestPacket(self, name, mode):
        return self.makeRequestPacket(WRQ, name, mode)

    def makeReadRequestPacket(self, name, mode):
        return self.makeRequestPacket(RRQ, name, mode)

    def makeRequestPacket(self, opcode, name, mode):
        lst = [pack(">H", opcode), name, "\0", mode, "\0"]
        return string.join(lst, "")

    def makeDataPacket(self, data):
        return pack(">HH", DATA, self.blockNumber) + data

    def makeAckPacket(self):
        return pack(">HH", ACK, self.blockNumber)

    def getNextPacket(self):
        now = time.time()
        if self.needResend:
            self.needResend = 0
            if self.trace:
                self.dumpPacket("sent", self.lastPacket)
            self.lastResendTime = now
            return self.lastPacket
        if self.needAck:
            return None
        packet = self.getProtocolPacket()
        self.lastPacket = packet
        self.lastPacketTime = now
        self.lastResendTime = now
        self.needAck = 1
        if self.trace:
            self.dumpPacket("sent", packet)
        return packet

    def handlePacket(self, packet, addr):
        if packet is None:
            self.handleTimeout()
            return
        (host, port) = addr
        if self.serverPort is None:
            self.serverPort = port
        elif port != self.serverPort:
            # Ignore packet sent from wrong UDP port.  This can happen
            # if the initial request packet is retransmitted because
            # of a timeout.  The RFC says that it is supposed to reply
            # with an ERROR packet back to that UDP port.  The BSD
            # version accepts the packet and uses that UDP port as the
            # new destination for sending packets.
            return
        if self.trace:
            self.dumpPacket("received", packet)
        self.handleProtocolPacket(packet, addr)

    def getTimeout(self):
        if self.done:
            return 0
        now = time.time()
        nextTime = min(
            self.lastPacketTime + self.maxTimeout,
            self.lastResendTime + self.resendTimeout)
        return max(0, nextTime - now)

    def handleTimeout(self):
        now = time.time()
        if now >= self.lastPacketTime + self.maxTimeout:
            self.log("timeout")
            self.done = 1
            raise Error("TFTP timeout")
        else:
            self.needResend = 1

    def getServerPortNumber(self):
        if self.serverPort is None:
            return self.tftpInitialPort
        return self.serverPort

    def isDone(self):
        return self.done

    def dumpPacket(self, where, data):
        limit = 16
        text = string.join(map(lambda x: "%02X" % ord(x), data[:limit]), " ")
        if len(data) > limit:
            text = text + ("... (%d bytes)" % len(data))
        self.log("%-10s %s" % (where, text))

    def log(self, text):
        print text

    def handleErrorPacket(self, packet):
        code = 0
        if len(packet) >= 4:
            (code,) = unpack(">H", packet[2:4])
        message = string.replace(packet[4:], "\0", "")
        if self.trace:
            self.log("Error code %d: %s" % (code, message))
        self.done = 1
        raise ProtocolError(code, message)


class TftpProtocolDownloader(TftpBase):

    def __init__(self, name, **options):
        apply(TftpBase.__init__, (self, name), options)

    def getProtocolPacket(self):
        if self.state == "start":
            packet = self.makeReadRequestPacket(self.name, self.mode)
            self.state = "data"
        elif self.state == "data":
            packet = self.makeAckPacket()
        self.blockNumber = self.blockNumber + 1
        if self.lastSegment:
            self.done = 1
        return packet

    def handleProtocolPacket(self, packet, addr):
        if len(packet) < 2:
            return
        (rxOpcode,) = unpack(">H", packet[:2])
        if rxOpcode == ERROR:
            self.handleErrorPacket(packet)
            return
        if rxOpcode != DATA or len(packet) < 4:
            return
        (rxBlock,) = unpack(">H", packet[2:4])
        data = packet[4:]
        if rxBlock != self.blockNumber & 0xffff:
            # On an error, try to synchronize
            # both sides.
            j = self.flushPacketInputQueue()
            if j > 0 and self.trace:
                self.log("discarded %d packets" % j)
            if rxBlock == (self.blockNumber - 1) & 0xffff:
                self.needResend = 1
            return
        self.writeData(data)
        self.needAck = 0
        if len(data) < self.segmentSize:
            self.lastSegment = 1


class TftpProtocolUploader(TftpBase):

    def __init__(self, name, **options):
        apply(TftpBase.__init__, (self, name), options)

    def getProtocolPacket(self):
        if self.state == "start":
            packet = self.makeWriteRequestPacket(self.name, self.mode)
            self.state = "data"
        elif self.state == "data":
            data = self.readData(self.segmentSize)
            if len(data) < self.segmentSize:
                self.lastSegment = 1
            packet = self.makeDataPacket(data)
        return packet

    def handleProtocolPacket(self, packet, addr):
        if len(packet) < 2:
            return
        (rxOpcode,) = unpack(">H", packet[:2])
        if rxOpcode == ERROR:
            self.handleErrorPacket(packet)
            return
        if rxOpcode != ACK or len(packet) < 4:
            return
        (rxBlock,) = unpack(">H", packet[2:4])
        if rxBlock != self.blockNumber & 0xffff:
            # On an error, try to synchronize
            # both sides.
            j = self.flushPacketInputQueue()
            if j > 0 and self.trace:
                self.log("discarded %d packets" % j)
            if rxBlock == (self.blockNumber - 1) & 0xffff:
                # Resend data in response to a duplicate ACK.  The BSD
                # version does this, but it really should not.  (See
                # RFC-1350, RFC-1123)
                self.needResend = 1
            return
        self.blockNumber = self.blockNumber + 1
        self.needAck = 0
        if self.lastSegment:
            self.done = 1


class SocketIO:

    def __init__(self, remoteHost):
        self.sendAddr = (remoteHost, 0)

    def sendPacket(self, packet):
        self.sock.sendto(packet, self.sendAddr)

    def recvPacket(self, timeout):
        now = time.time()
        deadline = now + timeout
        while 1:
            (sr, sw, se) = select.select([self.sock], [], [], deadline - now)
            if self.sock not in sr:
                return (None, None)
            (packet, addr) = self.sock.recvfrom(MAX_PACKET_SIZE)
            if len(packet) > 0:
                return (packet, addr)
            now = time.time()

    def transferLoop(self):
        while not self.isDone():
            packet = self.getNextPacket()
            if packet:
                self.sendAddr = (self.sendAddr[0], self.getServerPortNumber())
                self.sendPacket(packet)
            timeout = self.getTimeout()
            (packet, addr) = self.recvPacket(timeout)
            self.handlePacket(packet, addr)

    def transfer(self):
        self.sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        self.sock.bind(("", 0))
        self.sock.setblocking(0)
        self.transferLoop()

    def flushPacketInputQueue(self):
        # Attempt to clear any packets that have been received by the
        # OS but not by this application.  Of course we have no
        # control over any packets that might still be out there in a
        # router between the server and the client.
        j = 0
        while 1:
            buf = pack("i", 0)
            buf = fcntl.ioctl(self.sock.fileno(), termios.FIONREAD, buf)
            (i,) = unpack("i", buf)
            if i == 0:
                break
            j = j + 1
            junk = self.sock.recvfrom(MAX_PACKET_SIZE)
        return j


class TftpDownloader(SocketIO, TftpProtocolDownloader):

    def __init__(self, file, remoteHost, remoteFilename, **options):
        SocketIO.__init__(self, remoteHost)
        apply(TftpProtocolDownloader.__init__, (self, remoteFilename), options)
        self.file = file

    def writeData(self, data):
        return self.file.write(data)

class TftpUploader(SocketIO, TftpProtocolUploader):

    def __init__(self, file, remoteHost, remoteFilename, **options):
        SocketIO.__init__(self, remoteHost)
        apply(TftpProtocolUploader.__init__, (self, remoteFilename), options)
        self.file = file

    def readData(self, length):
        return self.file.read(length)


def _transfer(klass, local, host, remote, options, openMode):
    file = open(local, openMode)
    t = apply(klass, (file, host, remote), options)
    t.transfer()

def get(localFilename, remoteHost, remoteFilename, **options):
    _transfer(
        TftpDownloader, localFilename, remoteHost, remoteFilename,
        options, "wb")

def put(localFilename, remoteHost, remoteFilename, **options):
    _transfer(
        TftpUploader, localFilename, remoteHost, remoteFilename,
        options, "rb")


def tftpcp(argv):
    import getopt

    usage = """\
Usage: tftpcp [options] host:source dest
   or: tftpcp [options] source host:dest
Transfer files to/from a remote host using the TFTP protocol.
Options:
    -d                  Enable packet trace mode
    -t <seconds>        Set the transfer timeout
    -x <seconds>        Set the retransmit timeout
"""

    try:
        (optlist, args) = getopt.getopt(argv, "dt:x:")
    except getopt.error, ex:
        sys.stderr.write("%s\n" % ex)
        sys.stderr.write(usage)
        return 1

    if len(args) != 2:
        sys.stderr.write("Wrong number of arguments.\n")
        sys.stderr.write(usage)
        return 1

    options = {}
    for (k, v) in optlist:
        try:
            if k == "-d":
                options["trace"] = 1
            elif k == "-t":
                x = float(v)
                if not (0 < x <= 300):
                    raise ValueError("Valid range is 0 < x <= 300")
                options["maxTimeout"] = x
            elif k == "-x":
                x = float(v)
                if not (0 < x <= 300):
                    raise ValueError("Valid range is 0 < x <= 300")
                options["resendTimeout"] = x
        except ValueError, ex:
            sys.stderr.write(
                "The value %s is not valid for the %s option.\n(%s)\n" % (
                    repr(v), k, ex))
            sys.stderr.write(usage)
            return 1

    (src, dest) = map(lambda x: string.split(x, ":", 1), args)

    srcPath = src[-1]
    srcHost = len(src) > 1 and src[0] or ""
    destPath = dest[-1]
    destHost = len(dest) > 1 and dest[0] or ""

    if (srcHost and destHost) or (not srcHost and not destHost):
        sys.stderr.write(usage)
        return 1

    try:
        if srcHost:
            apply(get, (destPath, srcHost, srcPath), options)
        else:
            apply(put, (srcPath, destHost, destPath), options)
    except ProtocolError, (code, message):
        sys.stderr.write("TFTP Error code %d: %s\n" % (code, message))
        return 1
    except (Error, IOError), ex:
        sys.stderr.write("%s\n" % ex)
        return 1
    except socket.error, (code, message):
        sys.stderr.write("[Errno %d] %s\n" % (code, message))
        return 1

    return 0


if __name__ == "__main__":
    sys.exit(tftpcp(sys.argv[1:]))