1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
|
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)
|