#!/usr/bin/env python
#
#  Copyright (c) 1999, 2000, 2001 Sean Reifschneider, tummy.com, ltd.
#  All Rights Reserved

'''Functions for manipulation of net strings.

Netstrings (as described in ftp://koobera.math.uic.edu/www/proto/netstrings.txt)
are strings of the form "<length in bytes>:<string>,".  They allow the
programmer to allocate an appropriately sized buffer ahead of time, and also
to ensure that an entire string has been received (instead of relying on
parsing strings on CR-LF markers).

Simple netstring example:

	from netstring import *

	print 'Doing hello+there "%s"' % str(nstos('5:hello,5:there,', 1))
	print 'Doing empty+there "%s"' % str(nstos('0:,5:there,', 1))
	print 'Creating hello world!: "%s"' % stons('hello world!')
	print 'Iscomplete:', iscomplete('12:hello world')
	print 'Iscomplete:', iscomplete('12:hello world!')
	print 'Iscomplete:', iscomplete('12:hello world!,')
	print 'Iscomplete:', iscomplete('12:hello world!,12:hello world!,')

	s = stons(stons('name') + stons('value'))
	print 'Created embedded: "%s"' % s
	print '  Contains: "%s"' % nstos(s)
	s2 = nstos(s)
	while 1:
		if not s2: break
		s2, rest = nstos(s2, 1)
		print '    s2: "%s"' % s2
		print '      rest: "%s"' % rest
		s2 = rest
'''

revision = '$Revision: 1.14 $'

import string
import socket, select
import types
import os


class LengthOverflow:
	'''Exception class thrown when a netstring is larger than the
	user requested it be.'''
	def __init__(self, s): self.data = s
	def __str__(self): return(self.data)


def nstos(s, returnRemainder = 0):
	'''Convert a netstring into a regular string (possibly returning remainder).

	RETURNS: Parse netstring into regular string and return it.  If
		"returnRemainder" is true, a tuple of ( string, remainder ).
	EXCEPTIONS: ValueError is raised if argument is not a valid netstring.
	ARGUMENTS:
		- s -- string to parse into a regular string.
		- returnRemainder -- If true, a tuple is returned consisting of
			the parsed netstring, and the unparsed portion of the string.
			The remainder is empty if the entire string was consumed.
	'''
	colonpos = string.find(s, ':')
	if colonpos < 1: raise ValueError, 'Invalid format for netstring'
	try: size = int(s[:colonpos])
	except: raise ValueError, 'Left hand side of ":" must be only numeric'

	strstart = colonpos + 1
	termpos = strstart + size
	if len(s) - 1 < termpos:
		raise ValueError, 'Expected string of length %d, got %d' % ( termpos + 1,
				len(s) )
	if s[termpos] != ',':
		raise ValueError, 'Netstring must be terminated by a comma.'

	rawstr = s[strstart:termpos]
	if returnRemainder:
		remainder = s[termpos + 1:]
		return(( rawstr, remainder ))
	else:
		return(rawstr)

def stons(s):
	'''Convert a regular string into a netstring.

	RETURNS: Return argument converted into a netstring.
	EXCEPTIONS: None generated internally.
	ARGUMENTS:
		- s -- string to convert into a netstring.  If 's' is a list/tuple, each
		       element in turn is converted to a netstring, and the results are
				 concatenated and turned into a netstring.
	'''
	if type(s) == types.ListType:
		s2 = ''
		for s in s: s2 = s2 + stons(s)
		return(stons(s2))
	if type(s) == types.TupleType:
		s2 = ''
		for s in s: s2 = s2 + stons(s)
		return(stons(s2))

	return('%d:%s,' % ( len(s), s ))


################
def nstolist(s):
	'''Convert a netstring into a list.
	Given the netstring created by stons() when passed a list, this returns
	the original list.

	RETURNS: Return argument converted into a list.
	EXCEPTIONS: None generated internally.
	ARGUMENTS:
		- s -- netstring to convert to a list.
	'''
	s = nstos(s)
	list = []
	while s:
		e, s = nstos(s, returnRemainder = 1)
		list.append(e)
	return(list)


###################
def dicttons(dict):
	'''Convert a dictionary into a netstring.
	Given a dictionary, convert it into a netstring consisting of
	netstrings for each key/value pair, where each is converted
	using "str()".

	RETURNS: Argument converted into a netstring.
	EXCEPTIONS: None generated internally.
	ARGUMENTS:
		- dict -- dictionary to convert into a netstring.
	'''
	s = ''
	keys = dict.keys()
	keys.sort()
	for key in keys:
		s = s + stons(str(key)) + stons(str(dict[key]))
	return(stons(s))


#####################################
def nstodict(s, returnRemainder = 0):
	'''Convert a dictionary into a netstring.
	Given the netstring created by dicttons(), this converts it back into
	a dictionary.  All keys and values are strings, no matter their
	original type in the dictionary.

	RETURNS: Argument converted into a dictionary.
	EXCEPTIONS: ValueError -- If string doesn't contain an even number of
		sub-netstrings.
	ARGUMENTS:
		- s -- netstring to convert into a dictionary.
	'''
	dict = {}

	remainder = ''
	if returnRemainder:
		s, remainder = nstos(s, returnRemainder = 1)
	else:
		s = nstos(s)

	while s:
		try:
			key, s = nstos(s, returnRemainder = 1)
			value, s = nstos(s, returnRemainder = 1)
		except ValueError:
			raise ValueError, 'Argument must have an even number ' \
					'of sub-netstrings.'
		dict[key] = value

	if returnRemainder:
		return(( dict, remainder ))
	else:
		return(dict)


##################
def iscomplete(s):
	'''Return 1 if there are enough bytes to construct a netstring.

	RETURNS: 1 if there are enough bytes in argument to construct a netstring.
		For example if argument is "12:hello world!" (or shorter) this function
		will return 0, while if it's "12:hello world!," (or longer) it will
		return 1.
	EXCEPTIONS: None generated internally.
	ARGUMENTS:
		- s -- string to check to for completeness.
	'''
	colonpos = string.find(s, ':')
	if colonpos < 1: return 0
	try: size = int(s[:colonpos])
	except: raise ValueError, 'Left hand side of ":" must be only numeric'
	return(len(s) > size + colonpos + 1)


################################
def sockread(sock, bytesToRead):
	data = ''
	while bytesToRead > 0:
		rfdl = select.select([sock] , [], [], None)[0]
		if len(rfdl) == 0: return(data)
		s = rfdl[0].recv(bytesToRead)
		if not s: raise IOError, 'Error reading netstring data.'
		bytesToRead = bytesToRead - len(s)
		data = data + s
	return(data)


#############
class Reader:
	#######################################
	def __init__(self, src, maxLen = None):
		self.src = src
		if hasattr(self.src, 'read'): self._get = self.src.read
		elif hasattr(self.src, 'recv'): self._get = self.src.recv
		elif type(self.src) == types.IntType:
			self._get = lambda n, fd = self.src: os.read(fd, n)
		else: raise TypeError, 'Argument must have recv() or read() methods.'
		self.maxLen = maxLen
		self.len = 0
		self.startnew()


	##################################
	def startnew(self, maxLen = None):
		if not maxLen: maxLen = self.maxLen
		data = ''

		#  consume existing data in current netstring
		while self.len > 0: self.read(4096)

		self.len = 0
		while 1:
			s = self._get(1)
			if not s:
				raise ValueError, 'Ran out of data reading netstring header.'
			if s == ':':
				self.len = int(data)
				break
			if not s in string.digits:
				raise ValueError, 'Input stream had invalid character "%s"' % s
			data = data + s
			if maxLen and data and int(data) > maxLen or len(data) > maxLen:
				raise LengthOverflow, 'Netstring larger than requested limit.'


	##########################
	def read(self, toread = None):
		if toread == None or toread > self.len: toread = self.len
		if toread < 1: return('')
		data = self._get(toread)
		self.len = self.len - len(data)
		if self.len == 0:
			s = self._get(1)
			if s != ',':
				raise ValueError, 'Found character "%s" instead of ","' % s
		return(data)
	

	############################
	def __getattr__(self, attr):
		return(getattr(self.src, attr))


#class netstringSocket:
#	def __init__(self, sock):
#		self.sock = sock
#
#	def __getattr__(self, name):
#		attr = getattr(self.sock, name)
#		setattr(self, name, attr)
#		return(attr)
#
#socket(family, type, proto = 0):
#	return(netstringSocket(socket.socket(family, type, proto)))
#
#fromfd(fd, family, type, proto = 0):
#	return(netstringSocket(socket.fromfd(fd, family, type, proto)))

if __name__ == '__main__':
	s = stons(stons('name') + stons('value'))
	print 'Created embedded: "%s"' % s
	print '  Contains: "%s"' % nstos(s)
	s2 = nstos(s)
	while 1:
		if not s2: break
		s2, rest = nstos(s2, 1)
		print '    s2: "%s"' % s2
		print '      rest: "%s"' % rest
		s2 = rest
	print 'Doing hello+there "%s"' % str(nstos('5:hello,5:there,', 1))
	print 'Doing empty+there "%s"' % str(nstos('0:,5:there,', 1))
	print 'Creating hello world!: "%s"' % stons('hello world!')
	print 'Iscomplete:', iscomplete('12:hello world')
	print 'Iscomplete:', iscomplete('12:hello world!')
	print 'Iscomplete:', iscomplete('12:hello world!,')
	print 'Iscomplete:', iscomplete('12:hello world!,12:hello world!,')
	s = dicttons({ 'a' : 1, 'b' : 'abc', 'c' : 123 })
	print 'dicttons:', s
	print 'nstodict:', nstodict(s)