from tftp.protocol import TFTP
from twisted.internet.defer import inlineCallbacks, returnValue


class TFTPProtocol(TFTP):
    """Simple hack to stop sending packets, if the connection was closed."""

    def __init__(self, backend, _clock=None):
        TFTP.__init__(self, backend, _clock)
        self._real_session_sendData = None

    def _session_sendData(self, session, bytes):
        if session.transport.connected:
            # Send data, if "connection" is established
            self._real_session_sendData(bytes)
        elif session.timeout_watchdog is not None and \
             session.timeout_watchdog.active():
            # Kill the timeout watchdog, so that we do not get:
            # "Timed out after a successful transfer"
            session.timeout_watchdog.cancel()
    
    @inlineCallbacks
    def _startSession(self, datagram, addr, mode):
        session = yield TFTP._startSession(self, datagram, addr, mode)

        # monkey patch the sendData method
        self._real_session_sendData = session.session.sendData
        session.session.sendData = \
            lambda bytes: self._session_sendData(session.session, bytes)

        returnValue(session)