ttsService.py 5.7 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122
  1. import sys
  2. import os
  3. import rpyc
  4. from plumbum import cli
  5. from rpyc.utils.server import ThreadedServer, ForkingServer, OneShotServer
  6. from rpyc.utils.classic import DEFAULT_SERVER_PORT, DEFAULT_SERVER_SSL_PORT
  7. from rpyc.utils.registry import REGISTRY_PORT
  8. from rpyc.utils.registry import UDPRegistryClient, TCPRegistryClient
  9. from rpyc.utils.authenticators import SSLAuthenticator
  10. from rpyc.lib import setup_logger
  11. from rpyc.core import SlaveService
  12. class ClassicServer(cli.Application):
  13. mode = cli.SwitchAttr(["-m", "--mode"], cli.Set("threaded", "forking", "stdio", "oneshot"),
  14. default="threaded", help="The serving mode (threaded, forking, or 'stdio' for "
  15. "inetd, etc.)")
  16. port = cli.SwitchAttr(["-p", "--port"], cli.Range(0, 65535), default=None,
  17. help="The TCP listener port ("
  18. "default = {DEFAULT_SERVER_PORT!r}, "
  19. "default for SSL = {DEFAULT_SERVER_SSL_PORT!r})",
  20. group="Socket Options")
  21. host = cli.SwitchAttr(["--host"], str, default="", help="The host to bind to. "
  22. "The default is localhost", group="Socket Options")
  23. ipv6 = cli.Flag(["--ipv6"], help="Enable IPv6", group="Socket Options")
  24. logfile = cli.SwitchAttr("--logfile", str, default=None, help="Specify the log file to use; "
  25. "the default is stderr", group="Logging")
  26. quiet = cli.Flag(["-q", "--quiet"], help="Quiet mode (only errors will be logged)",
  27. group="Logging")
  28. ssl_keyfile = cli.SwitchAttr("--ssl-keyfile", cli.ExistingFile,
  29. help="The keyfile to use for SSL. Required for SSL", group="SSL",
  30. requires=["--ssl-certfile"])
  31. ssl_certfile = cli.SwitchAttr("--ssl-certfile", cli.ExistingFile,
  32. help="The certificate file to use for SSL. Required for SSL", group="SSL",
  33. requires=["--ssl-keyfile"])
  34. ssl_cafile = cli.SwitchAttr("--ssl-cafile", cli.ExistingFile,
  35. help="The certificate authority chain file to use for SSL. "
  36. "Optional; enables client-side authentication",
  37. group="SSL", requires=["--ssl-keyfile"])
  38. auto_register = cli.Flag("--register", help="Asks the server to attempt registering with "
  39. "a registry server. By default, the server will not attempt to register",
  40. group="Registry")
  41. registry_type = cli.SwitchAttr("--registry-type", cli.Set("UDP", "TCP"),
  42. default="UDP", help="Specify a UDP or TCP registry", group="Registry")
  43. registry_port = cli.SwitchAttr("--registry-port", cli.Range(0, 65535), default=REGISTRY_PORT,
  44. help="The registry's UDP/TCP port", group="Registry")
  45. registry_host = cli.SwitchAttr("--registry-host", str, default=None,
  46. help="The registry host machine. For UDP, the default is 255.255.255.255; "
  47. "for TCP, a value is required", group="Registry")
  48. def main(self):
  49. if not self.host:
  50. self.host = "::1" if self.ipv6 else "127.0.0.1"
  51. if self.registry_type == "UDP":
  52. if self.registry_host is None:
  53. self.registry_host = "255.255.255.255"
  54. self.registrar = UDPRegistryClient(ip=self.registry_host, port=self.registry_port)
  55. else:
  56. if self.registry_host is None:
  57. raise ValueError("With TCP registry, you must specify --registry-host")
  58. self.registrar = TCPRegistryClient(ip=self.registry_host, port=self.registry_port)
  59. if self.ssl_keyfile:
  60. self.authenticator = SSLAuthenticator(self.ssl_keyfile, self.ssl_certfile,
  61. self.ssl_cafile)
  62. default_port = DEFAULT_SERVER_SSL_PORT
  63. else:
  64. self.authenticator = None
  65. default_port = DEFAULT_SERVER_PORT
  66. if self.port is None:
  67. self.port = default_port
  68. setup_logger(self.quiet, self.logfile)
  69. if self.mode == "threaded":
  70. self._serve_mode(ThreadedServer)
  71. elif self.mode == "forking":
  72. self._serve_mode(ForkingServer)
  73. elif self.mode == "oneshot":
  74. self._serve_oneshot()
  75. elif self.mode == "stdio":
  76. self._serve_stdio()
  77. def _serve_mode(self, factory):
  78. t = factory(SlaveService, hostname=self.host, port=self.port,
  79. reuse_addr=True, ipv6=self.ipv6, authenticator=self.authenticator,
  80. registrar=self.registrar, auto_register=self.auto_register)
  81. t.start()
  82. def _serve_oneshot(self):
  83. t = OneShotServer(SlaveService, hostname=self.host, port=self.port,
  84. reuse_addr=True, ipv6=self.ipv6, authenticator=self.authenticator,
  85. registrar=self.registrar, auto_register=self.auto_register)
  86. t._listen()
  87. sys.stdout.write("rpyc-oneshot\n")
  88. sys.stdout.write(f"{t.host}\t{t.port}\n")
  89. sys.stdout.flush()
  90. t.start()
  91. def _serve_stdio(self):
  92. origstdin = sys.stdin
  93. origstdout = sys.stdout
  94. sys.stdin = open(os.devnull, "r")
  95. sys.stdout = open(os.devnull, "w")
  96. sys.stderr = open(os.devnull, "w")
  97. conn = rpyc.classic.connect_pipes(origstdin, origstdout)
  98. try:
  99. try:
  100. conn.serve_all()
  101. except KeyboardInterrupt:
  102. print("User interrupt!")
  103. finally:
  104. conn.close()
  105. if __name__ == "__main__":
  106. ClassicServer.run()