Demo/rpc/rpc.py

   1 # Sun RPC version 2 -- RFC1057.
   2
   3 # XXX There should be separate exceptions for the various reasons why
   4 # XXX an RPC can fail, rather than using RuntimeError for everything
   5
   6 # XXX The UDP version of the protocol resends requests when it does
   7 # XXX not receive a timely reply -- use only for idempotent calls!
   8
   9 # XXX There is no provision for call timeout on TCP connections
  10
  11 import xdr
  12 import socket
  13 import os
  14
  15 RPCVERSION = 2
  16
  17 CALL = 0
  18 REPLY = 1
  19
  20 AUTH_NULL = 0
  21 AUTH_UNIX = 1
  22 AUTH_SHORT = 2
  23 AUTH_DES = 3
  24
  25 MSG_ACCEPTED = 0
  26 MSG_DENIED = 1
  27
  28 SUCCESS = 0                             # RPC executed successfully
  29 PROG_UNAVAIL  = 1                       # remote hasn't exported program
  30 PROG_MISMATCH = 2                       # remote can't support version #
  31 PROC_UNAVAIL  = 3                       # program can't support procedure
  32 GARBAGE_ARGS  = 4                       # procedure can't decode params
  33
  34 RPC_MISMATCH = 0                        # RPC version number != 2
  35 AUTH_ERROR = 1                          # remote can't authenticate caller
  36
  37 AUTH_BADCRED      = 1                   # bad credentials (seal broken)
  38 AUTH_REJECTEDCRED = 2                   # client must begin new session
  39 AUTH_BADVERF      = 3                   # bad verifier (seal broken)
  40 AUTH_REJECTEDVERF = 4                   # verifier expired or replayed
  41 AUTH_TOOWEAK      = 5                   # rejected for security reasons
  42
  43
  44 class Packer(xdr.Packer):
  45
  46         def pack_auth(self, auth):
  47                 flavor, stuff = auth
  48                 self.pack_enum(flavor)
  49                 self.pack_opaque(stuff)
  50
  51         def pack_auth_unix(self, stamp, machinename, uid, gid, gids):
  52                 self.pack_uint(stamp)
  53                 self.pack_string(machinename)
  54                 self.pack_uint(uid)
  55                 self.pack_uint(gid)
  56                 self.pack_uint(len(gids))
  57                 for i in gids:
  58                         self.pack_uint(i)
  59
  60         def pack_callheader(self, xid, prog, vers, proc, cred, verf):
  61                 self.pack_uint(xid)
  62                 self.pack_enum(CALL)
  63                 self.pack_uint(RPCVERSION)
  64                 self.pack_uint(prog)
  65                 self.pack_uint(vers)
  66                 self.pack_uint(proc)
  67                 self.pack_auth(cred)
  68                 self.pack_auth(verf)
  69                 # Caller must add procedure-specific part of call
  70
  71         def pack_replyheader(self, xid, verf):
  72                 self.pack_uint(xid)
  73                 self.pack_enum(REPLY)
  74                 self.pack_uint(MSG_ACCEPTED)
  75                 self.pack_auth(verf)
  76                 self.pack_enum(SUCCESS)
  77                 # Caller must add procedure-specific part of reply
  78
  79
  80 # Exceptions
  81 BadRPCFormat = 'rpc.BadRPCFormat'
  82 BadRPCVersion = 'rpc.BadRPCVersion'
  83 GarbageArgs = 'rpc.GarbageArgs'
  84
  85 class Unpacker(xdr.Unpacker):
  86
  87         def unpack_auth(self):
  88                 flavor = self.unpack_enum()
  89                 stuff = self.unpack_opaque()
  90                 return (flavor, stuff)
  91
  92         def unpack_callheader(self):
  93                 xid = self.unpack_uint(xid)
  94                 temp = self.unpack_enum()
  95                 if temp <> CALL:
  96                         raise BadRPCFormat, 'no CALL but ' + `temp`
  97                 temp = self.unpack_uint()
  98                 if temp <> RPCVERSION:
  99                         raise BadRPCVerspion, 'bad RPC version ' + `temp`
 100                 prog = self.unpack_uint()
 101                 vers = self.unpack_uint()
 102                 proc = self.unpack_uint()
 103                 cred = self.unpack_auth()
 104                 verf = self.unpack_auth()
 105                 return xid, prog, vers, proc, cred, verf
 106                 # Caller must add procedure-specific part of call
 107
 108         def unpack_replyheader(self):
 109                 xid = self.unpack_uint()
 110                 mtype = self.unpack_enum()
 111                 if mtype <> REPLY:
 112                         raise RuntimeError, 'no REPLY but ' + `mtype`
 113                 stat = self.unpack_enum()
 114                 if stat == MSG_DENIED:
 115                         stat = self.unpack_enum()
 116                         if stat == RPC_MISMATCH:
 117                                 low = self.unpack_uint()
 118                                 high = self.unpack_uint()
 119                                 raise RuntimeError, \
 120                                   'MSG_DENIED: RPC_MISMATCH: ' + `low, high`
 121                         if stat == AUTH_ERROR:
 122                                 stat = self.unpack_uint()
 123                                 raise RuntimeError, \
 124                                         'MSG_DENIED: AUTH_ERROR: ' + `stat`
 125                         raise RuntimeError, 'MSG_DENIED: ' + `stat`
 126                 if stat <> MSG_ACCEPTED:
 127                         raise RuntimeError, \
 128                           'Neither MSG_DENIED nor MSG_ACCEPTED: ' + `stat`
 129                 verf = self.unpack_auth()
 130                 stat = self.unpack_enum()
 131                 if stat == PROG_UNAVAIL:
 132                         raise RuntimeError, 'call failed: PROG_UNAVAIL'
 133                 if stat == PROG_MISMATCH:
 134                         low = self.unpack_uint()
 135                         high = self.unpack_uint()
 136                         raise RuntimeError, \
 137                                 'call failed: PROG_MISMATCH: ' + `low, high`
 138                 if stat == PROC_UNAVAIL:
 139                         raise RuntimeError, 'call failed: PROC_UNAVAIL'
 140                 if stat == GARBAGE_ARGS:
 141                         raise RuntimeError, 'call failed: GARBAGE_ARGS'
 142                 if stat <> SUCCESS:
 143                         raise RuntimeError, 'call failed: ' + `stat`
 144                 return xid, verf
 145                 # Caller must get procedure-specific part of reply
 146
 147
 148 # Subroutines to create opaque authentication objects
 149
 150 def make_auth_null():
 151         return ''
 152
 153 def make_auth_unix(seed, host, uid, gid, groups):
 154         p = Packer()
 155         p.pack_auth_unix(seed, host, uid, gid, groups)
 156         return p.get_buf()
 157
 158 def make_auth_unix_default():
 159         try:
 160                 from os import getuid, getgid
 161                 uid = getuid()
 162                 gid = getgid()
 163         except ImportError:
 164                 uid = gid = 0
 165         import time
 166         return make_auth_unix(int(time.time()-unix_epoch()), \
 167                   socket.gethostname(), uid, gid, [])
 168
 169 _unix_epoch = -1
 170 def unix_epoch():
 171     """Very painful calculation of when the Unix Epoch is.
 172
 173     This is defined as the return value of time.time() on Jan 1st,
 174     1970, 00:00:00 GMT.
 175
 176     On a Unix system, this should always return 0.0.  On a Mac, the
 177     calculations are needed -- and hard because of integer overflow
 178     and other limitations.
 179
 180     """
 181     global _unix_epoch
 182     if _unix_epoch >= 0: return _unix_epoch
 183     import time
 184     now = time.time()
 185     localt = time.localtime(now)        # (y, m, d, hh, mm, ss, ..., ..., ...)
 186     gmt = time.gmtime(now)
 187     offset = time.mktime(localt) - time.mktime(gmt)
 188     y, m, d, hh, mm, ss = 1970, 1, 1, 0, 0, 0
 189     offset, ss = divmod(ss + offset, 60)
 190     offset, mm = divmod(mm + offset, 60)
 191     offset, hh = divmod(hh + offset, 24)
 192     d = d + offset
 193     _unix_epoch = time.mktime((y, m, d, hh, mm, ss, 0, 0, 0))
 194     print "Unix epoch:", time.ctime(_unix_epoch)
 195     return _unix_epoch
 196
 197
 198 # Common base class for clients
 199
 200 class Client:
 201
 202         def __init__(self, host, prog, vers, port):
 203                 self.host = host
 204                 self.prog = prog
 205                 self.vers = vers
 206                 self.port = port
 207                 self.makesocket() # Assigns to self.sock
 208                 self.bindsocket()
 209                 self.connsocket()
 210                 self.lastxid = 0 # XXX should be more random?
 211                 self.addpackers()
 212                 self.cred = None
 213                 self.verf = None
 214
 215         def close(self):
 216                 self.sock.close()
 217
 218         def makesocket(self):
 219                 # This MUST be overridden
 220                 raise RuntimeError, 'makesocket not defined'
 221
 222         def connsocket(self):
 223                 # Override this if you don't want/need a connection
 224                 self.sock.connect((self.host, self.port))
 225
 226         def bindsocket(self):
 227                 # Override this to bind to a different port (e.g. reserved)
 228                 self.sock.bind(('', 0))
 229
 230         def addpackers(self):
 231                 # Override this to use derived classes from Packer/Unpacker
 232                 self.packer = Packer()
 233                 self.unpacker = Unpacker('')
 234
 235         def make_call(self, proc, args, pack_func, unpack_func):
 236                 # Don't normally override this (but see Broadcast)
 237                 if pack_func is None and args is not None:
 238                         raise TypeError, 'non-null args with null pack_func'
 239                 self.start_call(proc)
 240                 if pack_func:
 241                         pack_func(args)
 242                 self.do_call()
 243                 if unpack_func:
 244                         result = unpack_func()
 245                 else:
 246                         result = None
 247                 self.unpacker.done()
 248                 return result
 249
 250         def start_call(self, proc):
 251                 # Don't override this
 252                 self.lastxid = xid = self.lastxid + 1
 253                 cred = self.mkcred()
 254                 verf = self.mkverf()
 255                 p = self.packer
 256                 p.reset()
 257                 p.pack_callheader(xid, self.prog, self.vers, proc, cred, verf)
 258
 259         def do_call(self):
 260                 # This MUST be overridden
 261                 raise RuntimeError, 'do_call not defined'
 262
 263         def mkcred(self):
 264                 # Override this to use more powerful credentials
 265                 if self.cred == None:
 266                         self.cred = (AUTH_NULL, make_auth_null())
 267                 return self.cred
 268
 269         def mkverf(self):
 270                 # Override this to use a more powerful verifier
 271                 if self.verf == None:
 272                         self.verf = (AUTH_NULL, make_auth_null())
 273                 return self.verf
 274
 275         def call_0(self):               # Procedure 0 is always like this
 276                 return self.make_call(0, None, None, None)
 277
 278
 279 # Record-Marking standard support
 280
 281 def sendfrag(sock, last, frag):
 282         x = len(frag)
 283         if last: x = x | 0x80000000L
 284         header = (chr(int(x>>24 & 0xff)) + chr(int(x>>16 & 0xff)) + \
 285                   chr(int(x>>8 & 0xff)) + chr(int(x & 0xff)))
 286         sock.send(header + frag)
 287
 288 def sendrecord(sock, record):
 289         sendfrag(sock, 1, record)
 290
 291 def recvfrag(sock):
 292         header = sock.recv(4)
 293         if len(header) < 4:
 294                 raise EOFError
 295         x = long(ord(header[0]))<<24 | ord(header[1])<<16 | \
 296             ord(header[2])<<8 | ord(header[3])
 297         last = ((x & 0x80000000) != 0)
 298         n = int(x & 0x7fffffff)
 299         frag = ''
 300         while n > 0:
 301                 buf = sock.recv(n)
 302                 if not buf: raise EOFError
 303                 n = n - len(buf)
 304                 frag = frag + buf
 305         return last, frag
 306
 307 def recvrecord(sock):
 308         record = ''
 309         last = 0
 310         while not last:
 311                 last, frag = recvfrag(sock)
 312                 record = record + frag
 313         return record
 314
 315
 316 # Try to bind to a reserved port (must be root)
 317
 318 last_resv_port_tried = None
 319 def bindresvport(sock, host):
 320         global last_resv_port_tried
 321         FIRST, LAST = 600, 1024 # Range of ports to try
 322         if last_resv_port_tried == None:
 323                 import os
 324                 last_resv_port_tried = FIRST + os.getpid() % (LAST-FIRST)
 325         for i in range(last_resv_port_tried, LAST) + \
 326                   range(FIRST, last_resv_port_tried):
 327                 last_resv_port_tried = i
 328                 try:
 329                         sock.bind((host, i))
 330                         return last_resv_port_tried
 331                 except socket.error, (errno, msg):
 332                         if errno <> 114:
 333                                 raise socket.error, (errno, msg)
 334         raise RuntimeError, 'can\'t assign reserved port'
 335
 336
 337 # Client using TCP to a specific port
 338
 339 class RawTCPClient(Client):
 340
 341         def makesocket(self):
 342                 self.sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
 343
 344         def do_call(self):
 345                 call = self.packer.get_buf()
 346                 sendrecord(self.sock, call)
 347                 reply = recvrecord(self.sock)
 348                 u = self.unpacker
 349                 u.reset(reply)
 350                 xid, verf = u.unpack_replyheader()
 351                 if xid <> self.lastxid:
 352                         # Can't really happen since this is TCP...
 353                         raise RuntimeError, 'wrong xid in reply ' + `xid` + \
 354                                 ' instead of ' + `self.lastxid`
 355
 356
 357 # Client using UDP to a specific port
 358
 359 class RawUDPClient(Client):
 360
 361         def makesocket(self):
 362                 self.sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
 363
 364         def do_call(self):
 365                 call = self.packer.get_buf()
 366                 self.sock.send(call)
 367                 try:
 368                         from select import select
 369                 except ImportError:
 370                         print 'WARNING: select not found, RPC may hang'
 371                         select = None
 372                 BUFSIZE = 8192 # Max UDP buffer size
 373                 timeout = 1
 374                 count = 5
 375                 while 1:
 376                         r, w, x = [self.sock], [], []
 377                         if select:
 378                                 r, w, x = select(r, w, x, timeout)
 379                         if self.sock not in r:
 380                                 count = count - 1
 381                                 if count < 0: raise RuntimeError, 'timeout'
 382                                 if timeout < 25: timeout = timeout *2
 383 ##                              print 'RESEND', timeout, count
 384                                 self.sock.send(call)
 385                                 continue
 386                         reply = self.sock.recv(BUFSIZE)
 387                         u = self.unpacker
 388                         u.reset(reply)
 389                         xid, verf = u.unpack_replyheader()
 390                         if xid <> self.lastxid:
 391 ##                              print 'BAD xid'
 392                                 continue
 393                         break
 394
 395
 396 # Client using UDP broadcast to a specific port
 397
 398 class RawBroadcastUDPClient(RawUDPClient):
 399
 400         def __init__(self, bcastaddr, prog, vers, port):
 401                 RawUDPClient.__init__(self, bcastaddr, prog, vers, port)
 402                 self.reply_handler = None
 403                 self.timeout = 30
 404
 405         def connsocket(self):
 406                 # Don't connect -- use sendto
 407                 self.sock.setsockopt(socket.SOL_SOCKET, socket.SO_BROADCAST, 1)
 408
 409         def set_reply_handler(self, reply_handler):
 410                 self.reply_handler = reply_handler
 411
 412         def set_timeout(self, timeout):
 413                 self.timeout = timeout # Use None for infinite timeout
 414
 415         def make_call(self, proc, args, pack_func, unpack_func):
 416                 if pack_func is None and args is not None:
 417                         raise TypeError, 'non-null args with null pack_func'
 418                 self.start_call(proc)
 419                 if pack_func:
 420                         pack_func(args)
 421                 call = self.packer.get_buf()
 422                 self.sock.sendto(call, (self.host, self.port))
 423                 try:
 424                         from select import select
 425                 except ImportError:
 426                         print 'WARNING: select not found, broadcast will hang'
 427                         select = None
 428                 BUFSIZE = 8192 # Max UDP buffer size (for reply)
 429                 replies = []
 430                 if unpack_func is None:
 431                         def dummy(): pass
 432                         unpack_func = dummy
 433                 while 1:
 434                         r, w, x = [self.sock], [], []
 435                         if select:
 436                                 if self.timeout is None:
 437                                         r, w, x = select(r, w, x)
 438                                 else:
 439                                         r, w, x = select(r, w, x, self.timeout)
 440                         if self.sock not in r:
 441                                 break
 442                         reply, fromaddr = self.sock.recvfrom(BUFSIZE)
 443                         u = self.unpacker
 444                         u.reset(reply)
 445                         xid, verf = u.unpack_replyheader()
 446                         if xid <> self.lastxid:
 447 ##                              print 'BAD xid'
 448                                 continue
 449                         reply = unpack_func()
 450                         self.unpacker.done()
 451                         replies.append((reply, fromaddr))
 452                         if self.reply_handler:
 453                                 self.reply_handler(reply, fromaddr)
 454                 return replies
 455
 456
 457 # Port mapper interface
 458
 459 # Program number, version and (fixed!) port number
 460 PMAP_PROG = 100000
 461 PMAP_VERS = 2
 462 PMAP_PORT = 111
 463
 464 # Procedure numbers
 465 PMAPPROC_NULL = 0                       # (void) -> void
 466 PMAPPROC_SET = 1                        # (mapping) -> bool
 467 PMAPPROC_UNSET = 2                      # (mapping) -> bool
 468 PMAPPROC_GETPORT = 3                    # (mapping) -> unsigned int
 469 PMAPPROC_DUMP = 4                       # (void) -> pmaplist
 470 PMAPPROC_CALLIT = 5                     # (call_args) -> call_result
 471
 472 # A mapping is (prog, vers, prot, port) and prot is one of:
 473
 474 IPPROTO_TCP = 6
 475 IPPROTO_UDP = 17
 476
 477 # A pmaplist is a variable-length list of mappings, as follows:
 478 # either (1, mapping, pmaplist) or (0).
 479
 480 # A call_args is (prog, vers, proc, args) where args is opaque;
 481 # a call_result is (port, res) where res is opaque.
 482
 483
 484 class PortMapperPacker(Packer):
 485
 486         def pack_mapping(self, mapping):
 487                 prog, vers, prot, port = mapping
 488                 self.pack_uint(prog)
 489                 self.pack_uint(vers)
 490                 self.pack_uint(prot)
 491                 self.pack_uint(port)
 492
 493         def pack_pmaplist(self, list):
 494                 self.pack_list(list, self.pack_mapping)
 495
 496         def pack_call_args(self, ca):
 497                 prog, vers, proc, args = ca
 498                 self.pack_uint(prog)
 499                 self.pack_uint(vers)
 500                 self.pack_uint(proc)
 501                 self.pack_opaque(args)
 502
 503
 504 class PortMapperUnpacker(Unpacker):
 505
 506         def unpack_mapping(self):
 507                 prog = self.unpack_uint()
 508                 vers = self.unpack_uint()
 509                 prot = self.unpack_uint()
 510                 port = self.unpack_uint()
 511                 return prog, vers, prot, port
 512
 513         def unpack_pmaplist(self):
 514                 return self.unpack_list(self.unpack_mapping)
 515
 516         def unpack_call_result(self):
 517                 port = self.unpack_uint()
 518                 res = self.unpack_opaque()
 519                 return port, res
 520
 521
 522 class PartialPortMapperClient:
 523
 524         def addpackers(self):
 525                 self.packer = PortMapperPacker()
 526                 self.unpacker = PortMapperUnpacker('')
 527
 528         def Set(self, mapping):
 529                 return self.make_call(PMAPPROC_SET, mapping, \
 530                         self.packer.pack_mapping, \
 531                         self.unpacker.unpack_uint)
 532
 533         def Unset(self, mapping):
 534                 return self.make_call(PMAPPROC_UNSET, mapping, \
 535                         self.packer.pack_mapping, \
 536                         self.unpacker.unpack_uint)
 537
 538         def Getport(self, mapping):
 539                 return self.make_call(PMAPPROC_GETPORT, mapping, \
 540                         self.packer.pack_mapping, \
 541                         self.unpacker.unpack_uint)
 542
 543         def Dump(self):
 544                 return self.make_call(PMAPPROC_DUMP, None, \
 545                         None, \
 546                         self.unpacker.unpack_pmaplist)
 547
 548         def Callit(self, ca):
 549                 return self.make_call(PMAPPROC_CALLIT, ca, \
 550                         self.packer.pack_call_args, \
 551                         self.unpacker.unpack_call_result)
 552
 553
 554 class TCPPortMapperClient(PartialPortMapperClient, RawTCPClient):
 555
 556         def __init__(self, host):
 557                 RawTCPClient.__init__(self, \
 558                         host, PMAP_PROG, PMAP_VERS, PMAP_PORT)
 559
 560
 561 class UDPPortMapperClient(PartialPortMapperClient, RawUDPClient):
 562
 563         def __init__(self, host):
 564                 RawUDPClient.__init__(self, \
 565                         host, PMAP_PROG, PMAP_VERS, PMAP_PORT)
 566
 567
 568 class BroadcastUDPPortMapperClient(PartialPortMapperClient, \
 569                                    RawBroadcastUDPClient):
 570
 571         def __init__(self, bcastaddr):
 572                 RawBroadcastUDPClient.__init__(self, \
 573                         bcastaddr, PMAP_PROG, PMAP_VERS, PMAP_PORT)
 574
 575
 576 # Generic clients that find their server through the Port mapper
 577
 578 class TCPClient(RawTCPClient):
 579
 580         def __init__(self, host, prog, vers):
 581                 pmap = TCPPortMapperClient(host)
 582                 port = pmap.Getport((prog, vers, IPPROTO_TCP, 0))
 583                 pmap.close()
 584                 if port == 0:
 585                         raise RuntimeError, 'program not registered'
 586                 RawTCPClient.__init__(self, host, prog, vers, port)
 587
 588
 589 class UDPClient(RawUDPClient):
 590
 591         def __init__(self, host, prog, vers):
 592                 pmap = UDPPortMapperClient(host)
 593                 port = pmap.Getport((prog, vers, IPPROTO_UDP, 0))
 594                 pmap.close()
 595                 if port == 0:
 596                         raise RuntimeError, 'program not registered'
 597                 RawUDPClient.__init__(self, host, prog, vers, port)
 598
 599
 600 class BroadcastUDPClient(Client):
 601
 602         def __init__(self, bcastaddr, prog, vers):
 603                 self.pmap = BroadcastUDPPortMapperClient(bcastaddr)
 604                 self.pmap.set_reply_handler(self.my_reply_handler)
 605                 self.prog = prog
 606                 self.vers = vers
 607                 self.user_reply_handler = None
 608                 self.addpackers()
 609
 610         def close(self):
 611                 self.pmap.close()
 612
 613         def set_reply_handler(self, reply_handler):
 614                 self.user_reply_handler = reply_handler
 615
 616         def set_timeout(self, timeout):
 617                 self.pmap.set_timeout(timeout)
 618
 619         def my_reply_handler(self, reply, fromaddr):
 620                 port, res = reply
 621                 self.unpacker.reset(res)
 622                 result = self.unpack_func()
 623                 self.unpacker.done()
 624                 self.replies.append((result, fromaddr))
 625                 if self.user_reply_handler is not None:
 626                         self.user_reply_handler(result, fromaddr)
 627
 628         def make_call(self, proc, args, pack_func, unpack_func):
 629                 self.packer.reset()
 630                 if pack_func:
 631                         pack_func(args)
 632                 if unpack_func is None:
 633                         def dummy(): pass
 634                         self.unpack_func = dummy
 635                 else:
 636                         self.unpack_func = unpack_func
 637                 self.replies = []
 638                 packed_args = self.packer.get_buf()
 639                 dummy_replies = self.pmap.Callit( \
 640                         (self.prog, self.vers, proc, packed_args))
 641                 return self.replies
 642
 643
 644 # Server classes
 645
 646 # These are not symmetric to the Client classes
 647 # XXX No attempt is made to provide authorization hooks yet
 648
 649 class Server:
 650
 651         def __init__(self, host, prog, vers, port):
 652                 self.host = host # Should normally be '' for default interface
 653                 self.prog = prog
 654                 self.vers = vers
 655                 self.port = port # Should normally be 0 for random port
 656                 self.makesocket() # Assigns to self.sock and self.prot
 657                 self.bindsocket()
 658                 self.host, self.port = self.sock.getsockname()
 659                 self.addpackers()
 660
 661         def register(self):
 662                 mapping = self.prog, self.vers, self.prot, self.port
 663                 p = TCPPortMapperClient(self.host)
 664                 if not p.Set(mapping):
 665                         raise RuntimeError, 'register failed'
 666
 667         def unregister(self):
 668                 mapping = self.prog, self.vers, self.prot, self.port
 669                 p = TCPPortMapperClient(self.host)
 670                 if not p.Unset(mapping):
 671                         raise RuntimeError, 'unregister failed'
 672
 673         def handle(self, call):
 674                 # Don't use unpack_header but parse the header piecewise
 675                 # XXX I have no idea if I am using the right error responses!
 676                 self.unpacker.reset(call)
 677                 self.packer.reset()
 678                 xid = self.unpacker.unpack_uint()
 679                 self.packer.pack_uint(xid)
 680                 temp = self.unpacker.unpack_enum()
 681                 if temp <> CALL:
 682                         return None # Not worthy of a reply
 683                 self.packer.pack_uint(REPLY)
 684                 temp = self.unpacker.unpack_uint()
 685                 if temp <> RPCVERSION:
 686                         self.packer.pack_uint(MSG_DENIED)
 687                         self.packer.pack_uint(RPC_MISMATCH)
 688                         self.packer.pack_uint(RPCVERSION)
 689                         self.packer.pack_uint(RPCVERSION)
 690                         return self.packer.get_buf()
 691                 self.packer.pack_uint(MSG_ACCEPTED)
 692                 self.packer.pack_auth((AUTH_NULL, make_auth_null()))
 693                 prog = self.unpacker.unpack_uint()
 694                 if prog <> self.prog:
 695                         self.packer.pack_uint(PROG_UNAVAIL)
 696                         return self.packer.get_buf()
 697                 vers = self.unpacker.unpack_uint()
 698                 if vers <> self.vers:
 699                         self.packer.pack_uint(PROG_MISMATCH)
 700                         self.packer.pack_uint(self.vers)
 701                         self.packer.pack_uint(self.vers)
 702                         return self.packer.get_buf()
 703                 proc = self.unpacker.unpack_uint()
 704                 methname = 'handle_' + `proc`
 705                 try:
 706                         meth = getattr(self, methname)
 707                 except AttributeError:
 708                         self.packer.pack_uint(PROC_UNAVAIL)
 709                         return self.packer.get_buf()
 710                 cred = self.unpacker.unpack_auth()
 711                 verf = self.unpacker.unpack_auth()
 712                 try:
 713                         meth() # Unpack args, call turn_around(), pack reply
 714                 except (EOFError, GarbageArgs):
 715                         # Too few or too many arguments
 716                         self.packer.reset()
 717                         self.packer.pack_uint(xid)
 718                         self.packer.pack_uint(REPLY)
 719                         self.packer.pack_uint(MSG_ACCEPTED)
 720                         self.packer.pack_auth((AUTH_NULL, make_auth_null()))
 721                         self.packer.pack_uint(GARBAGE_ARGS)
 722                 return self.packer.get_buf()
 723
 724         def turn_around(self):
 725                 try:
 726                         self.unpacker.done()
 727                 except RuntimeError:
 728                         raise GarbageArgs
 729                 self.packer.pack_uint(SUCCESS)
 730
 731         def handle_0(self): # Handle NULL message
 732                 self.turn_around()
 733
 734         def makesocket(self):
 735                 # This MUST be overridden
 736                 raise RuntimeError, 'makesocket not defined'
 737
 738         def bindsocket(self):
 739                 # Override this to bind to a different port (e.g. reserved)
 740                 self.sock.bind((self.host, self.port))
 741
 742         def addpackers(self):
 743                 # Override this to use derived classes from Packer/Unpacker
 744                 self.packer = Packer()
 745                 self.unpacker = Unpacker('')
 746
 747
 748 class TCPServer(Server):
 749
 750         def makesocket(self):
 751                 self.sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
 752                 self.prot = IPPROTO_TCP
 753
 754         def loop(self):
 755                 self.sock.listen(0)
 756                 while 1:
 757                         self.session(self.sock.accept())
 758
 759         def session(self, connection):
 760                 sock, (host, port) = connection
 761                 while 1:
 762                         try:
 763                                 call = recvrecord(sock)
 764                         except EOFError:
 765                                 break
 766                         except socket.error, msg:
 767                                 print 'socket error:', msg
 768                                 break
 769                         reply = self.handle(call)
 770                         if reply is not None:
 771                                 sendrecord(sock, reply)
 772
 773         def forkingloop(self):
 774                 # Like loop but uses forksession()
 775                 self.sock.listen(0)
 776                 while 1:
 777                         self.forksession(self.sock.accept())
 778
 779         def forksession(self, connection):
 780                 # Like session but forks off a subprocess
 781                 import os
 782                 # Wait for deceased children
 783                 try:
 784                         while 1:
 785                                 pid, sts = os.waitpid(0, 1)
 786                 except os.error:
 787                         pass
 788                 pid = None
 789                 try:
 790                         pid = os.fork()
 791                         if pid: # Parent
 792                                 connection[0].close()
 793                                 return
 794                         # Child
 795                         self.session(connection)
 796                 finally:
 797                         # Make sure we don't fall through in the parent
 798                         if pid == 0:
 799                                 os._exit(0)
 800
 801
 802 class UDPServer(Server):
 803
 804         def makesocket(self):
 805                 self.sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
 806                 self.prot = IPPROTO_UDP
 807
 808         def loop(self):
 809                 while 1:
 810                         self.session()
 811
 812         def session(self):
 813                 call, host_port = self.sock.recvfrom(8192)
 814                 reply = self.handle(call)
 815                 if reply <> None:
 816                         self.sock.sendto(reply, host_port)
 817
 818
 819 # Simple test program -- dump local portmapper status
 820
 821 def test():
 822         pmap = UDPPortMapperClient('')
 823         list = pmap.Dump()
 824         list.sort()
 825         for prog, vers, prot, port in list:
 826                 print prog, vers,
 827                 if prot == IPPROTO_TCP: print 'tcp',
 828                 elif prot == IPPROTO_UDP: print 'udp',
 829                 else: print prot,
 830                 print port
 831
 832
 833 # Test program for broadcast operation -- dump everybody's portmapper status
 834
 835 def testbcast():
 836         import sys
 837         if sys.argv[1:]:
 838                 bcastaddr = sys.argv[1]
 839         else:
 840                 bcastaddr = '<broadcast>'
 841         def rh(reply, fromaddr):
 842                 host, port = fromaddr
 843                 print host + '\t' + `reply`
 844         pmap = BroadcastUDPPortMapperClient(bcastaddr)
 845         pmap.set_reply_handler(rh)
 846         pmap.set_timeout(5)
 847         replies = pmap.Getport((100002, 1, IPPROTO_UDP, 0))
 848
 849
 850 # Test program for server, with corresponding client
 851 # On machine A: python -c 'import rpc; rpc.testsvr()'
 852 # On machine B: python -c 'import rpc; rpc.testclt()' A
 853 # (A may be == B)
 854
 855 def testsvr():
 856         # Simple test class -- proc 1 doubles its string argument as reply
 857         class S(UDPServer):
 858                 def handle_1(self):
 859                         arg = self.unpacker.unpack_string()
 860                         self.turn_around()
 861                         print 'RPC function 1 called, arg', `arg`
 862                         self.packer.pack_string(arg + arg)
 863         #
 864         s = S('', 0x20000000, 1, 0)
 865         try:
 866                 s.unregister()
 867         except RuntimeError, msg:
 868                 print 'RuntimeError:', msg, '(ignored)'
 869         s.register()
 870         print 'Service started...'
 871         try:
 872                 s.loop()
 873         finally:
 874                 s.unregister()
 875                 print 'Service interrupted.'
 876
 877
 878 def testclt():
 879         import sys
 880         if sys.argv[1:]: host = sys.argv[1]
 881         else: host = ''
 882         # Client for above server
 883         class C(UDPClient):
 884                 def call_1(self, arg):
 885                         return self.make_call(1, arg, \
 886                                 self.packer.pack_string, \
 887                                 self.unpacker.unpack_string)
 888         c = C(host, 0x20000000, 1)
 889         print 'making call...'
 890         reply = c.call_1('hello, world, ')
 891         print 'call returned', `reply`