Demo/rpc/rpc.py

# Implement (a subset of) Sun RPC, version 2 -- RFC1057.

# XXX There should be separate exceptions for the various reasons why
# XXX an RPC can fail, rather than using RuntimeError for everything

# XXX The UDP version of the protocol resends requests when it does
# XXX not receive a timely reply -- use only for idempotent calls!

import xdr
import socket
import os

RPCVERSION = 2

CALL = 0
REPLY = 1

AUTH_NULL = 0
AUTH_UNIX = 1
AUTH_SHORT = 2
AUTH_DES = 3

MSG_ACCEPTED = 0
MSG_DENIED = 1

SUCCESS = 0				# RPC executed successfully
PROG_UNAVAIL  = 1			# remote hasn't exported program
PROG_MISMATCH = 2			# remote can't support version #
PROC_UNAVAIL  = 3			# program can't support procedure
GARBAGE_ARGS  = 4			# procedure can't decode params

RPC_MISMATCH = 0			# RPC version number != 2
AUTH_ERROR = 1				# remote can't authenticate caller

AUTH_BADCRED      = 1			# bad credentials (seal broken)
AUTH_REJECTEDCRED = 2			# client must begin new session
AUTH_BADVERF      = 3			# bad verifier (seal broken)
AUTH_REJECTEDVERF = 4			# verifier expired or replayed
AUTH_TOOWEAK      = 5			# rejected for security reasons


class Packer(xdr.Packer):

	def pack_auth(self, auth):
		flavor, stuff = auth
		self.pack_enum(flavor)
		self.pack_opaque(stuff)

	def pack_auth_unix(self, stamp, machinename, uid, gid, gids):
		self.pack_uint(stamp)
		self.pack_string(machinename)
		self.pack_uint(uid)
		self.pack_uint(gid)
		self.pack_uint(len(gids))
		for i in gids:
			self.pack_uint(i)

	def pack_callheader(self, xid, prog, vers, proc, cred, verf):
		self.pack_uint(xid)
		self.pack_enum(CALL)
		self.pack_uint(RPCVERSION)
		self.pack_uint(prog)
		self.pack_uint(vers)
		self.pack_uint(proc)
		self.pack_auth(cred)
		self.pack_auth(verf)
		# Caller must add procedure-specific part of call

	def pack_replyheader(self, xid, verf):
		self.pack_uint(xid)
		self.pack_enum(REPLY)
		self.pack_uint(MSG_ACCEPTED)
		self.pack_auth(verf)
		self.pack_enum(SUCCESS)
		# Caller must add procedure-specific part of reply


# Exceptions
BadRPCFormat = 'rpc.BadRPCFormat'
BadRPCVersion = 'rpc.BadRPCVersion'
GarbageArgs = 'rpc.GarbageArgs'

class Unpacker(xdr.Unpacker):

	def unpack_auth(self):
		flavor = self.unpack_enum()
		stuff = self.unpack_opaque()
		return (flavor, stuff)

	def unpack_callheader(self):
		xid = self.unpack_uint(xid)
		temp = self.unpack_enum()
		if temp <> CALL:
			raise BadRPCFormat, 'no CALL but ' + `temp`
		temp = self.unpack_uint()
		if temp <> RPCVERSION:
			raise BadRPCVerspion, 'bad RPC version ' + `temp`
		prog = self.unpack_uint()
		vers = self.unpack_uint()
		proc = self.unpack_uint()
		cred = self.unpack_auth()
		verf = self.unpack_auth()
		return xid, prog, vers, proc, cred, verf
		# Caller must add procedure-specific part of call

	def unpack_replyheader(self):
		xid = self.unpack_uint()
		mtype = self.unpack_enum()
		if mtype <> REPLY:
			raise RuntimeError, 'no REPLY but ' + `mtype`
		stat = self.unpack_enum()
		if stat == MSG_DENIED:
			stat = self.unpack_enum()
			if stat == RPC_MISMATCH:
				low = self.unpack_uint()
				high = self.unpack_uint()
				raise RuntimeError, \
				  'MSG_DENIED: RPC_MISMATCH: ' + `low, high`
			if stat == AUTH_ERROR:
				stat = self.unpack_uint()
				raise RuntimeError, \
					'MSG_DENIED: AUTH_ERROR: ' + `stat`
			raise RuntimeError, 'MSG_DENIED: ' + `stat`
		if stat <> MSG_ACCEPTED:
			raise RuntimeError, \
			  'Neither MSG_DENIED nor MSG_ACCEPTED: ' + `stat`
		verf = self.unpack_auth()
		stat = self.unpack_enum()
		if stat == PROG_UNAVAIL:
			raise RuntimeError, 'call failed: PROG_UNAVAIL'
		if stat == PROG_MISMATCH:
			low = self.unpack_uint()
			high = self.unpack_uint()
			raise RuntimeError, \
				'call failed: PROG_MISMATCH: ' + `low, high`
		if stat == PROC_UNAVAIL:
			raise RuntimeError, 'call failed: PROC_UNAVAIL'
		if stat == GARBAGE_ARGS:
			raise RuntimeError, 'call failed: GARBAGE_ARGS'
		if stat <> SUCCESS:
			raise RuntimeError, 'call failed: ' + `stat`
		return xid, verf
		# Caller must get procedure-specific part of reply


# Subroutines to create opaque authentication objects

def make_auth_null():
	return ''

def make_auth_unix(seed, host, uid, gid, groups):
	p = Packer().init()
	p.pack_auth_unix(seed, host, uid, gid, groups)
	return p.get_buf()

def make_auth_unix_default():
	try:
		from os import getuid, getgid
		uid = getuid()
		gid = getgid()
	except ImportError:
		uid = gid = 0
	return make_auth_unix(0, socket.gethostname(), uid, gid, [])


# Common base class for clients

class Client:

	def init(self, host, prog, vers, port, type):
		self.host = host
		self.prog = prog
		self.vers = vers
		self.port = port
		self.type = type
		self.sock = socket.socket(socket.AF_INET, type)
		self.sock.connect((host, port))
		self.lastxid = 0
		self.addpackers()
		self.cred = None
		self.verf = None
		return self

	def Null(self):			# Procedure 0 is always like this
		self.start_call(0)
		self.do_call(0)
		self.end_call()

	def close(self):
		self.sock.close()

	# Functions that may be overridden by specific derived classes

	def addpackers(self):
		self.packer = Packer().init()
		self.unpacker = Unpacker().init('')

	def mkcred(self, proc):
		if self.cred == None:
			self.cred = (AUTH_NULL, make_auth_null())
		return self.cred

	def mkverf(self, proc):
		if self.verf == None:
			self.verf = (AUTH_NULL, make_auth_null())
		return self.verf


# Record-Marking standard support

def sendfrag(sock, last, frag):
	x = len(frag)
	if last: x = x | 0x80000000L
	header = (chr(int(x>>24 & 0xff)) + chr(int(x>>16 & 0xff)) + \
		  chr(int(x>>8 & 0xff)) + chr(int(x & 0xff)))
	sock.send(header + frag)

def sendrecord(sock, record):
	sendfrag(sock, 1, record)

def recvfrag(sock):
	header = sock.recv(4)
	if len(header) < 4:
		raise EOFError
	x = long(ord(header[0]))<<24 | ord(header[1])<<16 | \
	    ord(header[2])<<8 | ord(header[3])
	last = ((x & 0x80000000) != 0)
	n = int(x & 0x7fffffff)
	frag = ''
	while n > 0:
		buf = sock.recv(n)
		if not buf: raise EOFError
		n = n - len(buf)
		frag = frag + buf
	return last, frag

def recvrecord(sock):
	record = ''
	last = 0
	while not last:
		last, frag = recvfrag(sock)
		record = record + frag
	return record


# Raw TCP-based client

class RawTCPClient(Client):

	def init(self, host, prog, vers, port):
		return Client.init(self, host, prog, vers, port, \
			socket.SOCK_STREAM)

	def start_call(self, proc):
		self.lastxid = xid = self.lastxid + 1
		cred = self.mkcred(proc)
		verf = self.mkverf(proc)
		p = self.packer
		p.reset()
		p.pack_callheader(xid, self.prog, self.vers, proc, cred, verf)

	def do_call(self, *rest):
		# rest is used for UDP buffer size; ignored for TCP
		call = self.packer.get_buf()
		sendrecord(self.sock, call)
		reply = recvrecord(self.sock)
		u = self.unpacker
		u.reset(reply)
		xid, verf = u.unpack_replyheader()
		if xid <> self.lastxid:
			# Can't really happen since this is TCP...
			raise RuntimeError, 'wrong xid in reply ' + `xid` + \
				' instead of ' + `self.lastxid`

	def end_call(self):
		self.unpacker.done()


# Raw UDP-based client

class RawUDPClient(Client):

	def init(self, host, prog, vers, port):
		return Client.init(self, host, prog, vers, port, \
			socket.SOCK_DGRAM)

	def start_call(self, proc):
		self.lastxid = xid = self.lastxid + 1
		cred = self.mkcred(proc)
		verf = self.mkverf(proc)
		p = self.packer
		p.reset()
		p.pack_callheader(xid, self.prog, self.vers, proc, cred, verf)

	def do_call(self, *rest):
		try:
			from select import select
		except ImportError:
			print 'WARNING: select not found, RPC may hang'
			select = None
		if len(rest) == 0:
			bufsize = 8192
		elif len(rest) > 1:
			raise TypeError, 'too many args'
		else:
			bufsize = rest[0] + 512
		call = self.packer.get_buf()
		timeout = 1
		count = 5
		self.sock.send(call)
		while 1:
			r, w, x = [self.sock], [], []
			if select:
				r, w, x = select(r, w, x, timeout)
			if self.sock not in r:
				count = count - 1
				if count < 0: raise RuntimeError, 'timeout'
				if timeout < 25: timeout = timeout *2
				print 'RESEND', timeout, count
				self.sock.send(call)
				continue
			reply = self.sock.recv(bufsize)
			u = self.unpacker
			u.reset(reply)
			xid, verf = u.unpack_replyheader()
			if xid <> self.lastxid:
				print 'BAD xid'
				continue
			break

	def end_call(self):
		self.unpacker.done()


# Port mapper interface

PMAP_PORT = 111
PMAP_PROG = 100000
PMAP_VERS = 2
PMAPPROC_NULL = 0			# (void) -> void
PMAPPROC_SET = 1			# (mapping) -> bool
PMAPPROC_UNSET = 2			# (mapping) -> bool
PMAPPROC_GETPORT = 3			# (mapping) -> unsigned int
PMAPPROC_DUMP = 4			# (void) -> pmaplist
PMAPPROC_CALLIT = 5			# (call_args) -> call_result

# A mapping is (prog, vers, prot, port) and prot is one of:

IPPROTO_TCP = 6
IPPROTO_UDP = 17

# A pmaplist is a variable-length list of mappings, as follows:
# either (1, mapping, pmaplist) or (0).

# A call_args is (prog, vers, proc, args) where args is opaque;
# a call_result is (port, res) where res is opaque.


class PortMapperPacker(Packer):

	def pack_mapping(self, mapping):
		prog, vers, prot, port = mapping
		self.pack_uint(prog)
		self.pack_uint(vers)
		self.pack_uint(prot)
		self.pack_uint(port)

	def pack_pmaplist(self, list):
		self.pack_list(list, self.pack_mapping)


class PortMapperUnpacker(Unpacker):

	def unpack_mapping(self):
		prog = self.unpack_uint()
		vers = self.unpack_uint()
		prot = self.unpack_uint()
		port = self.unpack_uint()
		return prog, vers, prot, port

	def unpack_pmaplist(self):
		return self.unpack_list(self.unpack_mapping)


class PartialPortMapperClient:

	def addpackers(self):
		self.packer = PortMapperPacker().init()
		self.unpacker = PortMapperUnpacker().init('')

	def Set(self, mapping):
		self.start_call(PMAPPROC_SET)
		self.packer.pack_mapping(mapping)
		self.do_call()
		res = self.unpacker.unpack_uint()
		self.end_call()
		return res

	def Unset(self, mapping):
		self.start_call(PMAPPROC_UNSET)
		self.packer.pack_mapping(mapping)
		self.do_call()
		res = self.unpacker.unpack_uint()
		self.end_call()
		return res

	def Getport(self, mapping):
		self.start_call(PMAPPROC_GETPORT)
		self.packer.pack_mapping(mapping)
		self.do_call(4)
		port = self.unpacker.unpack_uint()
		self.end_call()
		return port

	def Dump(self):
		self.start_call(PMAPPROC_DUMP)
		self.do_call(8192-512)
		list = self.unpacker.unpack_pmaplist()
		self.end_call()
		return list


class TCPPortMapperClient(PartialPortMapperClient, RawTCPClient):

	def init(self, host):
		return RawTCPClient.init(self, \
			host, PMAP_PROG, PMAP_VERS, PMAP_PORT)


class UDPPortMapperClient(PartialPortMapperClient, RawUDPClient):

	def init(self, host):
		return RawUDPClient.init(self, \
			host, PMAP_PROG, PMAP_VERS, PMAP_PORT)


class TCPClient(RawTCPClient):

	def init(self, host, prog, vers):
		pmap = TCPPortMapperClient().init(host)
		port = pmap.Getport((prog, vers, IPPROTO_TCP, 0))
		if port == 0:
			raise RuntimeError, 'program not registered'
		pmap.close()
		return RawTCPClient.init(self, host, prog, vers, port)


class UDPClient(RawUDPClient):

	def init(self, host, prog, vers):
		pmap = UDPPortMapperClient().init(host)
		port = pmap.Getport((prog, vers, IPPROTO_UDP, 0))
		pmap.close()
		return RawUDPClient.init(self, host, prog, vers, port)


# Server classes

class Server:

	def init(self, host, prog, vers, port, type):
		self.host = host # Should normally be '' for default interface
		self.prog = prog
		self.vers = vers
		self.port = port # Should normally be 0 for random port
		self.type = type # SOCK_STREAM or SOCK_DGRAM
		self.sock = socket.socket(socket.AF_INET, type)
		self.sock.bind((host, port))
		self.host, self.port = self.sock.getsockname()
		self.addpackers()
		return self

	def register(self):
		if self.type == socket.SOCK_STREAM:
			type = IPPROTO_TCP
		elif self.type == socket.SOCK_DGRAM:
			type = IPPROTO_UDP
		else:
			raise ValueError, 'unknown protocol type'
		mapping = self.prog, self.vers, type, self.port
		p = TCPPortMapperClient().init(self.host)
		if not p.Set(mapping):
			raise RuntimeError, 'register failed'

	def unregister(self):
		if self.type == socket.SOCK_STREAM:
			type = IPPROTO_TCP
		elif self.type == socket.SOCK_DGRAM:
			type = IPPROTO_UDP
		else:
			raise ValueError, 'unknown protocol type'
		mapping = self.prog, self.vers, type, self.port
		p = TCPPortMapperClient().init(self.host)
		if not p.Unset(mapping):
			raise RuntimeError, 'unregister failed'

	def handle(self, call):
		# Don't use unpack_header but parse the header piecewise
		# XXX I have no idea if I am using the right error responses!
		self.unpacker.reset(call)
		self.packer.reset()
		xid = self.unpacker.unpack_uint()
		self.packer.pack_uint(xid)
		temp = self.unpacker.unpack_enum()
		if temp <> CALL:
			return None # Not worthy of a reply
		self.packer.pack_uint(REPLY)
		temp = self.unpacker.unpack_uint()
		if temp <> RPCVERSION:
			self.packer.pack_uint(MSG_DENIED)
			self.packer.pack_uint(RPC_MISMATCH)
			self.packer.pack_uint(RPCVERSION)
			self.packer.pack_uint(RPCVERSION)
			return self.packer.get_buf()
		self.packer.pack_uint(MSG_ACCEPTED)
		self.packer.pack_auth((AUTH_NULL, make_auth_null()))
		prog = self.unpacker.unpack_uint()
		if prog <> self.prog:
			self.packer.pack_uint(PROG_UNAVAIL)
			return self.packer.get_buf()
		vers = self.unpacker.unpack_uint()
		if vers <> self.vers:
			self.packer.pack_uint(PROG_MISMATCH)
			self.packer.pack_uint(self.vers)
			self.packer.pack_uint(self.vers)
			return self.packer.get_buf()
		proc = self.unpacker.unpack_uint()
		methname = 'handle_' + `proc`
		try:
			meth = getattr(self, methname)
		except AttributeError:
			self.packer.pack_uint(PROC_UNAVAIL)
			return self.packer.get_buf()
		cred = self.unpacker.unpack_auth()
		verf = self.unpacker.unpack_auth()
		try:
			meth() # Unpack args, call turn_around(), pack reply
		except (EOFError, GarbageArgs):
			# Too few or too many arguments
			self.packer.reset()
			self.packer.pack_uint(xid)
			self.packer.pack_uint(REPLY)
			self.packer.pack_uint(MSG_ACCEPTED)
			self.packer.pack_auth((AUTH_NULL, make_auth_null()))
			self.packer.pack_uint(GARBAGE_ARGS)
		return self.packer.get_buf()

	def turn_around(self):
		try:
			self.unpacker.done()
		except RuntimeError:
			raise GarbageArgs
		self.packer.pack_uint(SUCCESS)

	def handle_0(self): # Handle NULL message
		self.turn_around()

	# Functions that may be overridden by specific derived classes

	def addpackers(self):
		self.packer = Packer().init()
		self.unpacker = Unpacker().init('')


class TCPServer(Server):

	def init(self, host, prog, vers, port):
		return Server.init(self, host, prog, vers, port, \
			socket.SOCK_STREAM)

	def loop(self):
		self.sock.listen(0)
		while 1:
			self.session(self.sock.accept())

	def session(self, connection):
		sock, (host, port) = connection
		while 1:
			try:
				call = recvrecord(sock)
			except EOFError:
				break
			reply = self.handle(call)
			if reply <> None:
				sendrecord(sock, reply)


class UDPServer(Server):

	def init(self, host, prog, vers, port):
		return Server.init(self, host, prog, vers, port, \
			socket.SOCK_DGRAM)

	def loop(self):
		while 1:
			session()

	def session(self):
		call, host_port = self.sock.recvfrom(8192)
		reply = self.handle(call)
		if reply <> None:
			self.sock.sendto(reply, host_port)


# Simple test program -- dump local portmapper status

def test():
	import T
	T.TSTART()
	pmap = UDPPortMapperClient().init('')
	T.TSTOP()
	pmap.Null()
	T.TSTOP()
	list = pmap.Dump()
	T.TSTOP()
	list.sort()
	for prog, vers, prot, port in list:
		print prog, vers,
		if prot == IPPROTO_TCP: print 'tcp',
		elif prot == IPPROTO_UDP: print 'udp',
		else: print prot,
		print port


# Server and client test program.
# On machine A: python -c 'import rpc; rpc.testsvr()'
# On machine B: python -c 'import rpc; rpc.testclt()' A
# (A may be == B)

def testsvr():
	# Simple test class -- proc 1 doubles its string argument as reply
	class S(TCPServer):
		def handle_1(self):
			arg = self.unpacker.unpack_string()
			self.turn_around()
			print 'RPC function 1 called, arg', `arg`
			self.packer.pack_string(arg + arg)
	#
	s = S().init('', 0x20000000, 1, 0)
	try:
		s.unregister()
	except RuntimeError, msg:
		print 'RuntimeError:', msg, '(ignored)'
	s.register()
	print 'Service started...'
	try:
		s.loop()
	finally:
		s.unregister()
		print 'Service interrupted.'


def testclt():
	import sys
	if sys.argv[1:]: host = sys.argv[1]
	else: host = ''
	# Client for above server
	class C(TCPClient):
		def call_1(self, arg):
			self.start_call(1)
			self.packer.pack_string(arg)
			self.do_call()
			reply = self.unpacker.unpack_string()
			self.end_call()
			return reply
	c = C().init(host, 0x20000000, 1)
	print 'making call...'
	reply = c.call_1('hello, world, ')
	print 'call returned', `reply`