Major rewrite with new read_* interfaces

Demo/cwilib/telnetlib.py

-# Telnet client library
+# A TELNET client class.  Based on RFC 854: TELNET Protocol
+# Specification, by J. Postel and J. Reynolds
 
+
+# Example:
+#
+# >>> from telnetlib import Telnet
+# >>> tn = Telnet('voorn.cwi.nl', 79) # connect to finger port
+# >>> tn.write('guido\r\n')
+# >>> print tn.read_all()
+# Login name: guido                       In real life: Guido van Rossum
+# Office: M353,  x4127                    Home phone: 020-6225521
+# Directory: /ufs/guido                   Shell: /usr/local/bin/esh
+# On since Oct 28 11:02:16 on ttyq1   
+# Project: Multimedia Kernel Systems
+# No Plan.
+# >>>
+#
+# Note that read() won't read until eof -- it just reads some data
+# (but it guarantees to read at least one byte unless EOF is hit).
+#
+# It is possible to pass a Telnet object to select.select() in order
+# to wait until more data is available.  Note that in this case,
+# read_eager() may return '' even if there was data on the socket,
+# because the protocol negotiation may have eaten the data.
+# This is why EOFError is needed to distinguish between "no data"
+# and "connection closed" (since the socket also appears ready for
+# reading when it is closed).
+#
+# Bugs:
+# - may hang when connection is slow in the middle of an IAC sequence
+#
+# To do:
+# - option negotiation
+
+
+# Imported modules
 import socket
 import select
 import string
 import regsub
 
 # Tunable parameters
-TIMEOUT = 30.0
-DEBUGLEVEL = 1
+DEBUGLEVEL = 0
 
 # Telnet protocol defaults
 TELNET_PORT = 23
@@ -24,158 +58,264 @@ WILL = chr(251)
 
 class Telnet:
 
-	# Constructor
-	def __init__(self, host, port):
-		self.debuglevel = DEBUGLEVEL
-		self.host = host
-		if not port: port = TELNET_PORT
-		self.port = port
-		self.timeout = TIMEOUT
-		self.sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
-		self.sock.connect((self.host, self.port))
-		self.rawq = ''
-		self.irawq = 0
-		self.cookedq = ''
-
-	# Destructor
-	def __del__(self):
-		self.close()
-
-	# Print debug message
-	def msg(self, msg, *args):
-		if self.debuglevel > 0:
-			print 'TELNET:', msg%args
-
-	# Set debug level
-	def set_debuglevel(self, debuglevel):
-		self.debuglevel = debuglevel
-
-	# Set time-out on certain reads
-	def set_timeout(self, timeout):
-		self.timeout = float(timeout)
-
-	# Explicit close
-	def close(self):
-		if self.sock:
-			self.sock.close()
-		self.sock = None
-
-	# Return socket (e.g. for select)
-	def get_socket(self):
-		return self.sock
-
-	# Return socket's fileno (e.g. for select)
-	def fileno(self):
-		return self.sock.fileno()
-
-	# Write a string to the socket, doubling any IAC characters
-	def write(self, buffer):
-		if IAC in buffer:
-			buffer = regsub.gsub(IAC, IAC+IAC, buffer)
-		self.sock.send(buffer)
-
-	# Read until a given string is encountered or until timeout
-	def read_until(self, match):
-##		self.msg('read_until(%s)' % `match`)
-		n = len(match)
-		self.process_rawq()
-		i = string.find(self.cookedq, match)
-		if i < 0:
-			i = max(0, len(self.cookedq)-n)
-			self.fill_cookedq()
-			i = string.find(self.cookedq, match, i)
-		if i >= 0:
-			i = i+n
-			buf = self.cookedq[:i]
-			self.cookedq = self.cookedq[i:]
-##			self.msg('read_until(%s) -> %s' % (`match`, `buf`))
-			return buf
-		while select.select([self], [], [], self.timeout) == \
-			  ([self], [], []):
-			i = max(0, len(self.cookedq)-n)
-			self.fill_rawq()
-			self.process_rawq()
-			i = string.find(self.cookedq, match, i)
-			if i >= 0:
-				i = i+n
-				buf = self.cookedq[:i]
-				self.cookedq = self.cookedq[i:]
-##				self.msg('read_until(%s) -> %s' %
-##					  (`match`, `buf`))
-				return buf
-		buf = self.cookedq
-		self.cookedq = ''
-##		self.msg('read_until(%s) -> %s' % (`match`, `buf`))
-		return buf
+    # Constructor
+    def __init__(self, host, *args):
+	if not args:
+	    port = TELNET_PORT
+	else:
+	    if len(args) > 1: raise TypeError, 'too many args'
+	    port = args[0]
+	    if not port: port = TELNET_PORT
+	self.debuglevel = DEBUGLEVEL
+	self.host = host
+	self.port = port
+	self.sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+	self.sock.connect((self.host, self.port))
+	self.rawq = ''
+	self.irawq = 0
+	self.cookedq = ''
+	self.eof = 0
+
+    # Destructor
+    def __del__(self):
+	self.close()
+
+    # Debug message
+    def msg(self, msg, *args):
+	if self.debuglevel > 0:
+	    print 'Telnet(%s,%d):' % (self.host, self.port), msg % args
+
+    # Set debug level
+    def set_debuglevel(self, debuglevel):
+	self.debuglevel = debuglevel
+
+    # Explicit close
+    def close(self):
+	if self.sock:
+	    self.sock.close()
+	self.sock = None
+	self.eof = 1
+
+    # Return socket (e.g. for select)
+    def get_socket(self):
+	return self.sock
+
+    # Return socket's fileno (e.g. for select)
+    def fileno(self):
+	return self.sock.fileno()
+
+    # Write a string to the socket, doubling any IAC characters
+    # Might block if the connection is blocked
+    # May raise socket.error if the connection is closed
+    def write(self, buffer):
+	if IAC in buffer:
+	    buffer = regsub.gsub(IAC, IAC+IAC, buffer)
+	self.sock.send(buffer)
+
+    # The following read_* methods exist:
+    # Special case:
+    # - read_until() reads until a string is encountered or a timeout is hit
+    # These may block:
+    # - read_all() reads all data until EOF
+    # - read_some() reads at least one byte until EOF
+    # These may do I/O but won't block doing it:
+    # - read_very_eager() reads all data available on the socket
+    # - read_eager() reads either data already queued or some data
+    #                available on the socket
+    # These don't do I/O:
+    # - read_lazy() reads all data in the raw queue (processing it first)
+    # - read_very_lazy() reads all data in the cooked queue
 
-	# Read everything that's possible without really blocking
-	def read_now(self):
-		self.fill_cookedq()
-		buf = self.cookedq
-		self.cookedq = ''
-##		self.msg('read_now() --> %s' % `buf`)
+    # Read until a given string is encountered or until timeout
+    # Raise EOFError if connection closed and no cooked data available
+    # Return '' if no cooked data available otherwise
+    def read_until(self, match, *args):
+	if not args:
+	    timeout = None
+	else:
+	    if len(args) > 1: raise TypeError, 'too many args'
+	    timeout = args[0]
+	n = len(match)
+	self.process_rawq()
+	i = string.find(self.cookedq, match)
+	if i >= 0:
+	    i = i+n
+	    buf = self.cookedq[:i]
+	    self.cookedq = self.cookedq[i:]
+	    return buf
+	s_reply = ([self], [], [])
+	s_args = s_reply
+	if timeout is not None:
+	    s_args = s_args + (timeout,)
+	while not self.eof and apply(select.select, s_args) == s_reply:
+	    i = max(0, len(self.cookedq)-n)
+	    self.fill_rawq()
+	    self.process_rawq()
+	    i = string.find(self.cookedq, match, i)
+	    if i >= 0:
+		i = i+n
+		buf = self.cookedq[:i]
+		self.cookedq = self.cookedq[i:]
 		return buf
+	return self.read_very_lazy()
+
+    # Read all data until EOF
+    # Block until connection closed
+    def read_all(self):
+	self.process_rawq()
+	while not self.eof:
+	    self.fill_rawq()
+	    self.process_rawq()
+	buf = self.cookedq
+	self.cookedq = ''
+	return buf
+
+    # Read at least one byte of cooked data unless EOF is hit
+    # Return '' if EOF is hit
+    # Block if no data is immediately available
+    def read_some(self):
+	self.process_rawq()
+	while not self.cookedq and not self.eof:
+	    self.fill_rawq()
+	    self.process_rawq()
+	buf = self.cookedq
+	self.cookedq = ''
+	return buf
+
+    # Read everything that's possible without blocking in I/O (eager)
+    # Raise EOFError if connection closed and no cooked data available
+    # Return '' if no cooked data available otherwise
+    # Don't block unless in the midst of an IAC sequence
+    def read_very_eager(self):
+	self.process_rawq()
+	while not self.eof and self.sock_avail():
+	    self.fill_rawq()
+	    self.process_rawq()
+	return self.read_very_lazy()
+
+    # Read readily available data
+    # Raise EOFError if connection closed and no cooked data available
+    # Return '' if no cooked data available otherwise
+    # Don't block unless in the midst of an IAC sequence
+    def read_eager(self):
+	self.process_rawq()
+	while not self.cookedq and not self.eof and self.sock_avail():
+	    self.fill_rawq()
+	    self.process_rawq()
+	return self.read_very_lazy()
+
+    # Process and return data that's already in the queues (lazy)
+    # Raise EOFError if connection closed and no data available
+    # Return '' if no cooked data available otherwise
+    # Don't block unless in the midst of an IAC sequence
+    def read_lazy(self):
+	self.process_rawq()
+	return self.read_very_lazy()
+
+    # Return any data available in the cooked queue (very lazy)
+    # Raise EOFError if connection closed and no data available
+    # Return '' if no cooked data available otherwise
+    # Don't block
+    def read_very_lazy(self):
+	buf = self.cookedq
+	self.cookedq = ''
+	if not buf and self.eof and not self.rawq:
+	    raise EOFError, 'telnet connection closed'
+	return buf
+
+    # Transfer from raw queue to cooked queue
+    # Set self.eof when connection is closed
+    # Don't block unless in the midst of an IAC sequence
+    def process_rawq(self):
+	buf = ''
+	try:
+	    while self.rawq:
+		c = self.rawq_getchar()
+		if c != IAC:
+		    buf = buf + c
+		    continue
+		c = self.rawq_getchar()
+		if c == IAC:
+		    buf = buf + c
+		elif c in (DO, DONT):
+		    opt = self.rawq_getchar()
+		    self.msg('IAC %s %d', c == DO and 'DO' or 'DONT', ord(c))
+		    self.sock.send(IAC + WONT + opt)
+		elif c in (WILL, WONT):
+		    opt = self.rawq_getchar()
+		    self.msg('IAC %s %d',
+			  c == WILL and 'WILL' or 'WONT', ord(c))
+		else:
+		    self.msg('IAC %s not recognized' % `c`)
+	except EOFError: # raised by self.rawq_getchar()
+	    pass
+	self.cookedq = self.cookedq + buf
+
+    # Get next char from raw queue
+    # Block if no data is immediately available
+    # Raise EOFError when connection is closed
+    def rawq_getchar(self):
+	if not self.rawq:
+	    self.fill_rawq()
+	    if self.eof:
+		raise EOFError
+	c = self.rawq[self.irawq]
+	self.irawq = self.irawq + 1
+	if self.irawq >= len(self.rawq):
+	    self.rawq = ''
+	    self.irawq = 0
+	return c
+
+    # Fill raw queue from exactly one recv() system call
+    # Block if no data is immediately available
+    # Set self.eof when connection is closed
+    def fill_rawq(self):
+	if self.irawq >= len(self.rawq):
+	    self.rawq = ''
+	    self.irawq = 0
+	# The buffer size should be fairly small so as to avoid quadratic
+	# behavior in process_rawq() above
+	buf = self.sock.recv(50)
+	self.eof = (not buf)
+	self.rawq = self.rawq + buf
+
+    # Test whether data is available on the socket
+    def sock_avail(self):
+	return select.select([self], [], [], 0) == ([self], [], [])
+
 
-	# Fill cooked queue without blocking
-	def fill_cookedq(self):
-		self.process_rawq()
-		while select.select([self], [], [], 0) == ([self], [], []):
-			self.fill_rawq()
-			if not self.rawq:
-				raise EOFError
-			self.process_rawq()
-
-	# Transfer from raw queue to cooked queue
-	def process_rawq(self):
-		# There is some silliness going on here in an attempt
-		# to avoid quadratic behavior with large inputs...
-		buf = ''
-		while self.rawq:
-			c = self.rawq_getchar()
-			if c != IAC:
-				buf = buf + c
-				if len(buf) >= 44:
-##					self.msg('transfer: %s' % `buf`)
-					self.cookedq = self.cookedq + buf
-					buf = ''
-				continue
-			c = self.rawq_getchar()
-			if c == IAC:
-				buf = buf + c
-			elif c in (DO, DONT):
-				opt = self.rawq_getchar()
-				self.msg('IAC %s %d',
-					  c == DO and 'DO' or 'DONT',
-					  ord(c))
-				self.sock.send(IAC + WONT + opt)
-			elif c in (WILL, WONT):
-				opt = self.rawq_getchar()
-				self.msg('IAC %s %d',
-					  c == WILL and 'WILL' or 'WONT',
-					  ord(c))
-			else:
-				self.msg('IAC %s not recognized' % `c`)
-##		self.msg('transfer: %s' % `buf`)
-		self.cookedq = self.cookedq + buf
-
-	# Get next char from raw queue, blocking if necessary
-	def rawq_getchar(self):
-		if not self.rawq:
-			self.fill_rawq()
-		if self.irawq >= len(self.rawq):
-			raise EOFError
-		c = self.rawq[self.irawq]
-		self.irawq = self.irawq + 1
-		if self.irawq >= len(self.rawq):
-			self.rawq = ''
-			self.irawq = 0
-		return c
-
-	# Fill raw queue
-	def fill_rawq(self):
-		if self.irawq >= len(self.rawq):
-			self.rawq = ''
-			self.irawq = 0
-		buf = self.sock.recv(50)
-##		self.msg('fill_rawq(): %s' % `buf`)
-		self.rawq = self.rawq + buf
+# Test program
+# Usage: test [-d] ... [host [port]]
+def test():
+    import sys, string, socket, select
+    debuglevel = 0
+    while sys.argv[1:] and sys.argv[1] == '-d':
+	debuglevel = debuglevel+1
+	del sys.argv[1]
+    host = 'localhost'
+    if sys.argv[1:]:
+	host = sys.argv[1]
+    port = 0
+    if sys.argv[2:]:
+	portstr = sys.argv[2]
+	try:
+	    port = string.atoi(portstr)
+	except string.atoi_error:
+	    port = socket.getservbyname(portstr, 'tcp')
+    tn = Telnet(host, port)
+    tn.set_debuglevel(debuglevel)
+    while 1:
+	rfd, wfd, xfd = select.select([tn, sys.stdin], [], [])
+	if sys.stdin in rfd:
+	    line = sys.stdin.readline()
+	    tn.write(line)
+	if tn in rfd:
+	    try:
+		text = tn.read_eager()
+	    except EOFError:
+		print '*** Connection closed by remote host ***'
+		break
+	    if text:
+		sys.stdout.write(text)
+		sys.stdout.flush()
+    tn.close()