negtelnetserver.py 10 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349
  1. #!/usr/bin/env python
  2. # -*- coding: utf-8 -*-
  3. #
  4. """ A telnet server which negotiates"""
  5. from __future__ import (absolute_import, division, print_function,
  6. unicode_literals)
  7. import argparse
  8. import os
  9. import sys
  10. import logging
  11. import struct
  12. try: # Python 2
  13. import SocketServer as socketserver
  14. except ImportError: # Python 3
  15. import socketserver
  16. log = logging.getLogger(__name__)
  17. HOST = "localhost"
  18. IDENT = "NTEL"
  19. # The strings that indicate the test framework is checking our aliveness
  20. VERIFIED_REQ = b"verifiedserver"
  21. VERIFIED_RSP = b"WE ROOLZ: {pid}"
  22. def telnetserver(options):
  23. """
  24. Starts up a TCP server with a telnet handler and serves DICT requests
  25. forever.
  26. """
  27. if options.pidfile:
  28. pid = os.getpid()
  29. with open(options.pidfile, "w") as f:
  30. f.write(b"{0}".format(pid))
  31. local_bind = (HOST, options.port)
  32. log.info("Listening on %s", local_bind)
  33. # Need to set the allow_reuse on the class, not on the instance.
  34. socketserver.TCPServer.allow_reuse_address = True
  35. server = socketserver.TCPServer(local_bind, NegotiatingTelnetHandler)
  36. server.serve_forever()
  37. return ScriptRC.SUCCESS
  38. class NegotiatingTelnetHandler(socketserver.BaseRequestHandler):
  39. """Handler class for Telnet connections.
  40. """
  41. def handle(self):
  42. """
  43. Negotiates options before reading data.
  44. """
  45. neg = Negotiator(self.request)
  46. try:
  47. # Send some initial negotiations.
  48. neg.send_do("NEW_ENVIRON")
  49. neg.send_will("NEW_ENVIRON")
  50. neg.send_dont("NAWS")
  51. neg.send_wont("NAWS")
  52. # Get the data passed through the negotiator
  53. data = neg.recv(1024)
  54. log.debug("Incoming data: %r", data)
  55. if VERIFIED_REQ in data:
  56. log.debug("Received verification request from test framework")
  57. response_data = VERIFIED_RSP.format(pid=os.getpid())
  58. else:
  59. log.debug("Received normal request - echoing back")
  60. response_data = data.strip()
  61. if response_data:
  62. log.debug("Sending %r", response_data)
  63. self.request.sendall(response_data)
  64. except IOError:
  65. log.exception("IOError hit during request")
  66. class Negotiator(object):
  67. NO_NEG = 0
  68. START_NEG = 1
  69. WILL = 2
  70. WONT = 3
  71. DO = 4
  72. DONT = 5
  73. def __init__(self, tcp):
  74. self.tcp = tcp
  75. self.state = self.NO_NEG
  76. def recv(self, bytes):
  77. """
  78. Read bytes from TCP, handling negotiation sequences
  79. :param bytes: Number of bytes to read
  80. :return: a buffer of bytes
  81. """
  82. buffer = bytearray()
  83. # If we keep receiving negotiation sequences, we won't fill the buffer.
  84. # Keep looping while we can, and until we have something to give back
  85. # to the caller.
  86. while len(buffer) == 0:
  87. data = self.tcp.recv(bytes)
  88. if not data:
  89. # TCP failed to give us any data. Break out.
  90. break
  91. for byte in data:
  92. byte_int = self.byte_to_int(byte)
  93. if self.state == self.NO_NEG:
  94. self.no_neg(byte, byte_int, buffer)
  95. elif self.state == self.START_NEG:
  96. self.start_neg(byte_int)
  97. elif self.state in [self.WILL, self.WONT, self.DO, self.DONT]:
  98. self.handle_option(byte_int)
  99. else:
  100. # Received an unexpected byte. Stop negotiations
  101. log.error("Unexpected byte %s in state %s",
  102. byte_int,
  103. self.state)
  104. self.state = self.NO_NEG
  105. return buffer
  106. def byte_to_int(self, byte):
  107. return struct.unpack(b'B', byte)[0]
  108. def no_neg(self, byte, byte_int, buffer):
  109. # Not negotiating anything thus far. Check to see if we
  110. # should.
  111. if byte_int == NegTokens.IAC:
  112. # Start negotiation
  113. log.debug("Starting negotiation (IAC)")
  114. self.state = self.START_NEG
  115. else:
  116. # Just append the incoming byte to the buffer
  117. buffer.append(byte)
  118. def start_neg(self, byte_int):
  119. # In a negotiation.
  120. log.debug("In negotiation (%s)",
  121. NegTokens.from_val(byte_int))
  122. if byte_int == NegTokens.WILL:
  123. # Client is confirming they are willing to do an option
  124. log.debug("Client is willing")
  125. self.state = self.WILL
  126. elif byte_int == NegTokens.WONT:
  127. # Client is confirming they are unwilling to do an
  128. # option
  129. log.debug("Client is unwilling")
  130. self.state = self.WONT
  131. elif byte_int == NegTokens.DO:
  132. # Client is indicating they can do an option
  133. log.debug("Client can do")
  134. self.state = self.DO
  135. elif byte_int == NegTokens.DONT:
  136. # Client is indicating they can't do an option
  137. log.debug("Client can't do")
  138. self.state = self.DONT
  139. else:
  140. # Received an unexpected byte. Stop negotiations
  141. log.error("Unexpected byte %s in state %s",
  142. byte_int,
  143. self.state)
  144. self.state = self.NO_NEG
  145. def handle_option(self, byte_int):
  146. if byte_int in [NegOptions.BINARY,
  147. NegOptions.CHARSET,
  148. NegOptions.SUPPRESS_GO_AHEAD,
  149. NegOptions.NAWS,
  150. NegOptions.NEW_ENVIRON]:
  151. log.debug("Option: %s", NegOptions.from_val(byte_int))
  152. # No further negotiation of this option needed. Reset the state.
  153. self.state = self.NO_NEG
  154. else:
  155. # Received an unexpected byte. Stop negotiations
  156. log.error("Unexpected byte %s in state %s",
  157. byte_int,
  158. self.state)
  159. self.state = self.NO_NEG
  160. def send_message(self, message):
  161. packed_message = self.pack(message)
  162. self.tcp.sendall(packed_message)
  163. def pack(self, arr):
  164. return struct.pack(b'{0}B'.format(len(arr)), *arr)
  165. def send_iac(self, arr):
  166. message = [NegTokens.IAC]
  167. message.extend(arr)
  168. self.send_message(message)
  169. def send_do(self, option_str):
  170. log.debug("Sending DO %s", option_str)
  171. self.send_iac([NegTokens.DO, NegOptions.to_val(option_str)])
  172. def send_dont(self, option_str):
  173. log.debug("Sending DONT %s", option_str)
  174. self.send_iac([NegTokens.DONT, NegOptions.to_val(option_str)])
  175. def send_will(self, option_str):
  176. log.debug("Sending WILL %s", option_str)
  177. self.send_iac([NegTokens.WILL, NegOptions.to_val(option_str)])
  178. def send_wont(self, option_str):
  179. log.debug("Sending WONT %s", option_str)
  180. self.send_iac([NegTokens.WONT, NegOptions.to_val(option_str)])
  181. class NegBase(object):
  182. @classmethod
  183. def to_val(cls, name):
  184. return getattr(cls, name)
  185. @classmethod
  186. def from_val(cls, val):
  187. for k in cls.__dict__.keys():
  188. if getattr(cls, k) == val:
  189. return k
  190. return "<unknown>"
  191. class NegTokens(NegBase):
  192. # The start of a negotiation sequence
  193. IAC = 255
  194. # Confirm willingness to negotiate
  195. WILL = 251
  196. # Confirm unwillingness to negotiate
  197. WONT = 252
  198. # Indicate willingness to negotiate
  199. DO = 253
  200. # Indicate unwillingness to negotiate
  201. DONT = 254
  202. # The start of sub-negotiation options.
  203. SB = 250
  204. # The end of sub-negotiation options.
  205. SE = 240
  206. class NegOptions(NegBase):
  207. # Binary Transmission
  208. BINARY = 0
  209. # Suppress Go Ahead
  210. SUPPRESS_GO_AHEAD = 3
  211. # NAWS - width and height of client
  212. NAWS = 31
  213. # NEW-ENVIRON - environment variables on client
  214. NEW_ENVIRON = 39
  215. # Charset option
  216. CHARSET = 42
  217. def get_options():
  218. parser = argparse.ArgumentParser()
  219. parser.add_argument("--port", action="store", default=9019,
  220. type=int, help="port to listen on")
  221. parser.add_argument("--verbose", action="store", type=int, default=0,
  222. help="verbose output")
  223. parser.add_argument("--pidfile", action="store",
  224. help="file name for the PID")
  225. parser.add_argument("--logfile", action="store",
  226. help="file name for the log")
  227. parser.add_argument("--srcdir", action="store", help="test directory")
  228. parser.add_argument("--id", action="store", help="server ID")
  229. parser.add_argument("--ipv4", action="store_true", default=0,
  230. help="IPv4 flag")
  231. return parser.parse_args()
  232. def setup_logging(options):
  233. """
  234. Set up logging from the command line options
  235. """
  236. root_logger = logging.getLogger()
  237. add_stdout = False
  238. formatter = logging.Formatter("%(asctime)s %(levelname)-5.5s "
  239. "[{ident}] %(message)s"
  240. .format(ident=IDENT))
  241. # Write out to a logfile
  242. if options.logfile:
  243. handler = logging.FileHandler(options.logfile, mode="w")
  244. handler.setFormatter(formatter)
  245. handler.setLevel(logging.DEBUG)
  246. root_logger.addHandler(handler)
  247. else:
  248. # The logfile wasn't specified. Add a stdout logger.
  249. add_stdout = True
  250. if options.verbose:
  251. # Add a stdout logger as well in verbose mode
  252. root_logger.setLevel(logging.DEBUG)
  253. add_stdout = True
  254. else:
  255. root_logger.setLevel(logging.INFO)
  256. if add_stdout:
  257. stdout_handler = logging.StreamHandler(sys.stdout)
  258. stdout_handler.setFormatter(formatter)
  259. stdout_handler.setLevel(logging.DEBUG)
  260. root_logger.addHandler(stdout_handler)
  261. class ScriptRC(object):
  262. """Enum for script return codes"""
  263. SUCCESS = 0
  264. FAILURE = 1
  265. EXCEPTION = 2
  266. class ScriptException(Exception):
  267. pass
  268. if __name__ == '__main__':
  269. # Get the options from the user.
  270. options = get_options()
  271. # Setup logging using the user options
  272. setup_logging(options)
  273. # Run main script.
  274. try:
  275. rc = telnetserver(options)
  276. except Exception as e:
  277. log.exception(e)
  278. rc = ScriptRC.EXCEPTION
  279. log.info("Returning %d", rc)
  280. sys.exit(rc)