factory.py 6.7 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164
  1. """
  2. rpyc connection factories
  3. """
  4. import socket
  5. import thread, threading
  6. from rpyc import Connection, Channel, SocketStream, PipeStream, VoidService
  7. from rpyc.utils.registry import UDPRegistryClient
  8. from rpyc.lib import safe_import
  9. ssl = safe_import("ssl")
  10. class DiscoveryError(Exception):
  11. pass
  12. #------------------------------------------------------------------------------
  13. # API
  14. #------------------------------------------------------------------------------
  15. def connect_channel(channel, service = VoidService, config = {}):
  16. """creates a connection over a given channel
  17. channel - the channel to use
  18. service - the local service to expose (defaults to Void)
  19. config - configuration dict"""
  20. return Connection(service, channel, config = config)
  21. def connect_stream(stream, service = VoidService, config = {}):
  22. """creates a connection over a given stream
  23. stream - the stream to use
  24. service - the local service to expose (defaults to Void)
  25. config - configuration dict"""
  26. return connect_channel(Channel(stream), service = service, config = config)
  27. def connect_pipes(input, output, service = VoidService, config = {}):
  28. """creates a connection over the given input/output pipes
  29. input - the input pipe
  30. output - the output pipe
  31. service - the local service to expose (defaults to Void)
  32. config - configuration dict"""
  33. return connect_stream(PipeStream(input, output), service = service, config = config)
  34. def connect_stdpipes(service = VoidService, config = {}):
  35. """creates a connection over the standard input/output pipes
  36. service - the local service to expose (defaults to Void)
  37. config - configuration dict"""
  38. return connect_stream(PipeStream.from_std(), service = service, config = config)
  39. def connect(host, port, service = VoidService, config = {}):
  40. """creates a socket-connection to the given host
  41. host - the hostname to connect to
  42. port - the TCP port to use
  43. service - the local service to expose (defaults to Void)
  44. config - configuration dict"""
  45. return Connection(service, Channel(SocketStream.connect(host, port)), config = config)
  46. def tlslite_connect(host, port, username, password, service = VoidService, config = {}):
  47. """creates a TLS-connection to the given host (encrypted and authenticated)
  48. username - the username used to authenticate the client
  49. password - the password used to authenticate the client
  50. host - the hostname to connect to
  51. port - the TCP port to use
  52. service - the local service to expose (defaults to Void)
  53. config - configuration dict"""
  54. s = SocketStream.tlslite_connect(host, port, username, password)
  55. return Connection(service, Channel(s), config = config)
  56. def ssl_connect(host, port, keyfile = None, certfile = None, ca_certs = None,
  57. ssl_version = None, service = VoidService, config = {}):
  58. """creates an SSL-wrapped connection to the given host (encrypted and
  59. authenticated).
  60. host - the hostname to connect to
  61. port - the TCP port to use
  62. service - the local service to expose (defaults to Void)
  63. config - configuration dict
  64. keyfile, certfile, ca_certs, ssl_version -- arguments to ssl.wrap_socket.
  65. see that module's documentation for further info."""
  66. kwargs = {"server_side" : False}
  67. if keyfile:
  68. kwargs["keyfile"] = keyfile
  69. if certfile:
  70. kwargs["certfile"] = certfile
  71. if ca_certs:
  72. kwargs["ca_certs"] = ca_certs
  73. if ssl_version:
  74. kwargs["ssl_version"] = ssl_version
  75. else:
  76. kwargs["ssl_version"] = ssl.PROTOCOL_TLSv1
  77. s = SocketStream.ssl_connect(host, port, kwargs)
  78. return Connection(service, Channel(s), config = config)
  79. def discover(service_name, host = None, registrar = None, timeout = 2):
  80. """discovers hosts running the given service
  81. service_name - the service to look for
  82. host - limit the discovery to the given host only (None means any host)
  83. registrar - use this registry client to discover services. if None,
  84. use the default UDPRegistryClient with the default settings.
  85. timeout - the number of seconds to wait for a reply from the registry
  86. if no hosts are found, raises DiscoveryError
  87. returns a list of (ip, port) pairs
  88. """
  89. if registrar is None:
  90. registrar = UDPRegistryClient(timeout = timeout)
  91. addrs = registrar.discover(service_name)
  92. if not addrs:
  93. raise DiscoveryError("no servers exposing %r were found" % (service_name,))
  94. if host:
  95. ips = socket.gethostbyname_ex(host)[2]
  96. addrs = [(h, p) for h, p in addrs if h in ips]
  97. if not addrs:
  98. raise DiscoveryError("no servers exposing %r were found on %r" % (service_name, host))
  99. return addrs
  100. def connect_by_service(service_name, host = None, service = VoidService, config = {}):
  101. """create a connection to an arbitrary server that exposes the requested service
  102. service_name - the service to discover
  103. host - limit discovery to the given host only (None means any host)
  104. service - the local service to expose (defaults to Void)
  105. config - configuration dict"""
  106. host, port = discover(service_name, host = host)[0]
  107. return connect(host, port, service, config = config)
  108. def connect_subproc(args, service = VoidService, config = {}):
  109. """runs an rpyc server on a child process that and connects to it over
  110. the stdio pipes. uses the subprocess module.
  111. args - the args to Popen, e.g., ["python", "-u", "myfile.py"]
  112. service - the local service to expose (defaults to Void)
  113. config - configuration dict"""
  114. from subprocess import Popen, PIPE
  115. proc = Popen(args, stdin = PIPE, stdout = PIPE)
  116. conn = connect_pipes(proc.stdout, proc.stdin, service = service, config = config)
  117. conn.proc = proc # just so you can have control over the processs
  118. return conn
  119. def connect_thread(service = VoidService, config = {}, remote_service = VoidService, remote_config = {}):
  120. """starts an rpyc server on a thread and connects to it over a socket.
  121. service - the local service to expose (defaults to Void)
  122. config - configuration dict
  123. server_service - the remote service to expose (of the server; defaults to Void)
  124. server_config - remote configuration dict (of the server)
  125. """
  126. listener = socket.socket()
  127. listener.bind(("localhost", 0))
  128. listener.listen(1)
  129. def server(listener = listener):
  130. client = listener.accept()[0]
  131. listener.close()
  132. conn = connect_stream(SocketStream(client), service = remote_service,
  133. config = remote_config)
  134. try:
  135. conn.serve_all()
  136. except KeyboardInterrupt:
  137. thread.interrupt_main()
  138. t = threading.Thread(target = server)
  139. t.setDaemon(True)
  140. t.start()
  141. host, port = listener.getsockname()
  142. return connect(host, port, service = service, config = config)