import errno
import logging
import os
import struct

import SocketServer


logger = logging.getLogger(__name__)


class TFTPUDPHandler(SocketServer.BaseRequestHandler):
    """ Mixin. Implementation of the TFTP protocol.

    http://tools.ietf.org/html/rfc1350
    http://tools.ietf.org/html/rfc1782
    """

    # opcodes
    OP_RRQ = 1
    OP_WRQ = 2
    OP_DATA = 3
    OP_ACK = 4
    OP_ERROR = 5
    OP_OACK = 6

    # errors
    ERR_UNDEFINED = 0
    ERR_NOT_FOUND = 1
    ERR_PERM = 2
    ERR_DISK_FULL = 3
    ERR_ILLEGAL_OPERATION = 4
    ERR_UNKNOWN_TID = 5
    ERR_ALREADY_EXISTS = 6
    ERR_NO_SUCH_USER = 7

    def _log(self, level, msg, extra=None, exc_info=False):
        """ Add client_ip to extra.
        """
        log_extra = {'client_ip': self.client_address[0]}
        if extra:
            log_extra.update(extra)
        logger.log(level, msg, extra=log_extra, exc_info=exc_info)

    def handle(self):
        """ Extract header info and dispatch to handle_* methods.
        """
        data = self.request[0]
        opcode, = struct.unpack('!h', data[0:2])

        data = data[2:]  # skip opcode

        if opcode == self.OP_RRQ:
            args = data.split('\x00')

            if args[-1] != '':
                return self.send_error(
                    self.ERR_ILLEGAL_OPERATION,
                    'Final argument should end with a \\0'
                )
            args = args[:-1]  # skip final \0

            if len(args) < 2:
                return self.send_error(
                    self.ERR_ILLEGAL_OPERATION,
                    "Filename and mode required"
                )

            filename, mode = args[0], args[1]
            options = args[2:]

            # options must be a list like: [optname1, optvalue1, ..., optnameN,
            # optvalueN)
            if len(options) % 2:
                return self.send_error(
                    self.ERR_ILLEGAL_OPERATION,
                    "Malformed options"
                )

            # transform options to a dict
            options = dict(options[i:i+2] for i in xrange(0, len(options), 2))

            self.handle_rrq(filename, mode, options)

        elif opcode == self.OP_ACK:
            block_id, = struct.unpack('!h', data)
            self.handle_ack(block_id)

        else:
            return self.send_error(
                self.ERR_ILLEGAL_OPERATION,
                'Opcode %d not handled by the server' % opcode
            )

    def handle_rrq(self, filename, mode, options):
        """ Handle READ requests.

        Create a new session.
        """
        self._log(logging.INFO, 'GET %s (%s)' % (filename, mode))

        if mode.lower() != 'octet':
            self.send_error(
                self.ERR_ILLEGAL_OPERATION,
                'Only octet mode is supported by the server'
            )
            return

        try:
            filename = self.sanitize_filename(filename)
        except ValueError as exc:
            self.send_error(self.ERR_PERM, str(exc))
            return

        try:
            handle = self.load_file(filename)
        except IOError as exc:
            # If ENOENT, consider the file is missing. Otherwise, consider we
            # don't have the permission to read it.
            err_msg = exc.strerror or str(exc)
            if exc.errno == errno.ENOENT:
                self.send_error(
                    self.ERR_NOT_FOUND, '%s (%s)' % (err_msg, filename)
                )
            else:
                self.send_error(
                    self.ERR_PERM, '%s (%s)' % (err_msg, filename)
                )
            return
        # The file cannot be loaded for any (critical) reason. Log the
        # traceback.
        except Exception as exc:
            self._log(logging.ERROR, 'Internal error', exc_info=True)
            self.send_error(self.ERR_UNDEFINED, 'Internal error')
            return

        session =  TFTPSession(filename, handle)
        self.server.sessions[self.client_address] = session

        # If there is a supported option, return a OACK, otherwise return the
        # first packet.
        # For now, only 'blksize' is supported.
        blksize = options.get('blksize')

        # Set the block size and return a OACK
        if blksize:
            try:
                session.blksize = int(blksize)
            except ValueError:  # not an int
                return self.send_error(
                    self.ERR_ILLEGAL_OPERATION, 'Bad option value'
                )
            return self.send_oack(blksize=blksize)

        # No options, return the first part of the file
        self.send_data()

    def sanitize_filename(self, filename):
        """ Compute a filename that can be loaded by load_file. Make security
        checks to prevent path transversal.
        """
        raise NotImplementedError()

    def load_file(self, filename):
        """ Return a FILE like object.
        """
        raise NotImplementedError()

    def handle_ack(self, block_id):
        """ Client has aknowledged a block id. Can be a retransmission or the
        next packet to send.
        """
        self._log(logging.DEBUG, 'ACK (block %s)' % block_id)
        session = self.server.sessions.get(self.client_address)

        # If the ACK does not correspond to a read request
        if not session:
            return

        # Last packet was received
        if block_id == session.block_id + 1:

            # Final ACK from the client, kill the sesssion
            if session.last_read_is_eof:
                self._log(
                    logging.INFO,
                    'Transfer of %s successful' % session.filename
                )
                del self.server.sessions[self.client_address]
                return

            # Next packet
            session.block_id += 1

        # Send the next packet, or retransmit the last packet if there was an
        # error
        self.send_data()

    def send_oack(self, **options):
        """ Send options acknowledgement.
        """
        packed = struct.pack('!h', self.OP_OACK)
        for key, value in options.iteritems():
            packed += key + '\x00' + value + '\x00'

        socket = self.request[1]
        socket.sendto(packed, self.client_address)

    def send_data(self):
        """ Send the next data packet to the client.
        """
        session = self.server.sessions[self.client_address]
        session.handle.seek(session.block_id * session.blksize)
        data = session.handle.read(session.blksize)
        session.last_read_is_eof = len(data) < session.blksize

        packed = struct.pack('!hh', self.OP_DATA, session.block_id + 1)
        packed += data

        socket = self.request[1]
        socket.sendto(packed, self.client_address)

    def send_error(self, error_code, error_msg):
        """ Send error packet to the client.
        """
        self._log(logging.ERROR, error_msg)
        socket = self.request[1]
        packed = struct.pack('!hh', self.OP_ERROR, error_code)
        packed += error_msg + '\x00'
        socket.sendto(packed, self.client_address)


class TFTPSession(object):

    def __init__(self, filename, handle):
        self.block_id = 0
        self.filename = filename
        self.handle = handle
        self.last_read_is_eof = False
        self.blksize = 512


class FileSystemHandler(TFTPUDPHandler):

    def sanitize_filename(self, filename):
        """ Raise if trying to open a file up to the root folder.
        """
        server_root = os.path.abspath(self.server.root)
        abs_path = os.path.abspath(os.path.join(server_root, filename))
        if os.path.commonprefix([abs_path, server_root]) != server_root:
            raise ValueError('Directory trasversal prevented')
        return abs_path

    def load_file(self, filename):
        return open(filename)


class TFTPServer(SocketServer.ThreadingMixIn, SocketServer.UDPServer):

    def __init__(self, host='', port=69, root='/var/lib/tftpboot',
                 handler=FileSystemHandler):

        self.sessions = {}
        self.root = root
        SocketServer.UDPServer.__init__(self, (host, port), handler)
