def test_stream_sb(stream_parser): events = stream_parser.stream_updates( [ Tokenizer.Option(B.SB, None, 42), Tokenizer.StreamData(b"1234"), Tokenizer.Command(B.SE, B.SE.value), ] ) assert events == [StreamParser.OptionSubnegotiation(None, 42)]
def __init__(self, reader, writer): self.reader = reader self.writer = writer self.encoder = StreamStuffer() self.tokenizer = Tokenizer() self.parser = StreamParser() self.line_buffer = LineBuffer() self.update_buffer = [] self.prompt_mgr = None self.echo_state = EchoOptionState()
def test_stream_command_in_crlf(stream_parser): events = stream_parser.stream_updates( [ Tokenizer.StreamData(b"abc\r"), Tokenizer.Command(B.NOP, B.NOP.value), Tokenizer.StreamData(b"\ndef"), ] ) assert events == [ StreamParser.UserData("abc"), StreamParser.Command(B.NOP, B.NOP.value), StreamParser.UserData("\n"), StreamParser.UserData("def"), ]
def test_stream_user_data_crlf(stream_parser): events = stream_parser.stream_updates([Tokenizer.StreamData(b"Hello,\r\nworld!")]) assert events == [ StreamParser.UserData("Hello,"), StreamParser.UserData("\n"), StreamParser.UserData("world!"), ]
def test_stream_dont(stream_parser): events = stream_parser.stream_updates( [Tokenizer.Option(B.DONT, OPTIONS.TM, OPTIONS.TM.value)] ) assert events == [ StreamParser.OptionNegotiation( OPTIONS.TM, OPTIONS.TM.value, StreamParser.Host.LOCAL, False ) ]
def test_stream_user_data(stream_parser): events = stream_parser.stream_updates([Tokenizer.StreamData(b"Hello, world!")]) assert events == [StreamParser.UserData("Hello, world!")]
def test_split_option(tokenizer): toks = tokenizer.tokens(b"abc" + B.IAC.byte + B.WONT.byte) assert toks == [Tokenizer.StreamData(b"abc")] toks = tokenizer.tokens(bytes([42]) + b"def") assert toks == [Tokenizer.Option(B.WONT, None, 42), Tokenizer.StreamData(b"def")]
def test_split_command(tokenizer): toks = tokenizer.tokens(b"abc" + B.IAC.byte) assert toks == [Tokenizer.StreamData(b"abc")] toks = tokenizer.tokens(B.NOP.byte + b"def") assert toks == [Tokenizer.Command(B.NOP, B.NOP.value), Tokenizer.StreamData(b"def")]
def test_tokenizer_option_unrecognized(tokenizer): toks = tokenizer.tokens(B.IAC.byte + B.WILL.byte + bytes([42])) assert toks == [Tokenizer.Option(B.WILL.value, None, 42)]
def test_tokenizer_option_recognized(tokenizer): toks = tokenizer.tokens(B.IAC.byte + B.WILL.byte + OPTIONS.TM.byte) assert toks == [Tokenizer.Option(B.WILL.value, OPTIONS.TM, OPTIONS.TM.value)]
def test_stream_command(stream_parser): events = stream_parser.stream_updates([Tokenizer.Command(B.NOP, B.NOP.value)]) assert events == [StreamParser.Command(B.NOP, B.NOP.value)]
def tokenizer(): return Tokenizer()
def test_stream_wont(stream_parser): events = stream_parser.stream_updates([Tokenizer.Option(B.WONT, None, 42)]) assert events == [ StreamParser.OptionNegotiation(None, 42, StreamParser.Host.PEER, False) ]
def test_stream_will(stream_parser): events = stream_parser.stream_updates([Tokenizer.Option(B.WILL, None, 42)]) assert events == [ StreamParser.OptionNegotiation(None, 42, StreamParser.Host.PEER, True) ]
def test_tokenizer_data(tokenizer): TEST_DATA = b"abcde" toks = tokenizer.tokens(TEST_DATA) assert toks == [Tokenizer.StreamData(TEST_DATA)]
def test_stream_user_data_nonascii(stream_parser): events = stream_parser.stream_updates([Tokenizer.StreamData(b"abc\xabdef")]) assert events == [StreamParser.UserData("abcdef")]
def test_tokenizer_command(tokenizer): toks = tokenizer.tokens(B.IAC.byte + B.NOP.byte) assert toks == [Tokenizer.Command(B.NOP, B.NOP.value)]
def test_stream_iac(stream_parser): # IAC IAC produces a single IAC on the user stream. This is not a valid # ascii character, so it's filtered by decoding. events = stream_parser.stream_updates([Tokenizer.Command(B.IAC, B.IAC.value)]) assert events == []
class Terminal: READ_SIZE = 2**12 # arbitrary pleasant number? def __init__(self, reader, writer): self.reader = reader self.writer = writer self.encoder = StreamStuffer() self.tokenizer = Tokenizer() self.parser = StreamParser() self.line_buffer = LineBuffer() self.update_buffer = [] self.prompt_mgr = None self.echo_state = EchoOptionState() async def __aenter__(self): return self async def __aexit__(self, exc_type, exc_value, traceback): await self.close() async def close(self): await self.writer.drain() self.writer.close() await self.writer.wait_closed() async def sleep(self, secs): await self.writer.drain() await asyncio.sleep(secs) async def write(self, texts, drain=False): if texts is None: texts = [] if not isinstance(texts, list): texts = [texts] for text in texts: self._write(text) if drain: await self.writer.drain() def _write(self, text): if isinstance(text, str): text = StreamParser.UserData(text) out_data = self.encoder.stuff(text) logger.debug("writing", extra={"data": out_data}) self.writer.write(out_data) @contextlib.contextmanager def prompt(self, prompt): self.prompt_mgr = Prompt(prompt) try: yield finally: self.prompt_mgr = None async def input(self, prompt): with self.prompt(prompt): return await self._input_line() async def input_secret(self, prompt): with self.prompt(prompt): async with self.echo_off(): result = await self._input_line() # echo peer's LF await self.write("\n") return result async def _input_line(self): while True: line = await self._handle_input_once() if line: return line async def _handle_input_once(self): await self.require_line_buffer() annotations, line = self.line_buffer.pop() for annotation in annotations: handler_name = "annotation_" + annotation.__class__.__name__ handler = getattr(self, handler_name) await handler(annotation) return line async def require_line_buffer(self): while not self.line_buffer.has_line(): await self.require_update_buffer() update = self.update_buffer.pop(0) await self.write(self.handle_update(update), drain=True) async def require_update_buffer(self): while not self.update_buffer: await self.write(self.prompt_mgr.require_has_prompt(), drain=True) self.update_buffer = await self.fetch_updates() async def fetch_updates(self): data = await self.reader.read(self.READ_SIZE) logger.debug("read", extra={"data": data}) if not data: raise EOFError() toks = self.tokenizer.tokens(data) return self.parser.stream_updates(toks) async def annotation_TimingMark(self, annotation): await self.write(annotation.option.accept(), drain=True) # # locally requested state changes # @contextlib.asynccontextmanager async def echo_off(self): # The way we turn echo off is by turning on the server echo option # and then just not echoing anythiing. await self.write(self.echo_state.local_request(True), drain=True) try: yield finally: await self.write(self.echo_state.local_request(False), drain=True) # # input stream updates # def handle_update(self, update): handle = getattr(self, "update_" + update.__class__.__name__, None) if not handle: logger.info("unhandled update", extra={"update": update}) return return handle(update) def update_UserData(self, update): self.line_buffer.append(update.data) self.prompt_mgr.mark_user_data() # # telnet options # def update_OptionNegotiation(self, request): handler = self.get_option_handler(request) return handler(request) def get_option_handler(self, request): if request.option is None: return self.option_unhandled handler_name = "option_" + request.option.name handler = getattr(self, "option_" + request.option.name, None) return handler or self.option_unhandled def option_TM(self, request): if not request.state: # client is rejecting this option. we only send it when it's # requested, so it's safe to ignore this. logger.debug("ignoring unmatched request", extra={"request": request}) return if request.host == StreamParser.Host.LOCAL: # client is requesting a TM. logger.info(f"ACCEPTING", extra={"request": request}) self.line_buffer.annotate(self.TimingMark(request)) return # otherwise client is sending us a TM, which we didn't request. ignore it. logger.debug("ignoring unrequested request", extra={"request": request}) def option_ECHO(self, request): return self.echo_state.handle_negotiation(request) def option_unhandled(self, request): if request.state: # Peer is requesting an unsupported option. Refuse. logger.info("rejecting option", extra={"request": request}) return request.refuse() # else client is requesting to disable an option. we support no # options, so all are off. ignore it. else: logger.debug("ignoring option", extra={"request": request}) # telnet commands def update_Command(self, command): handler = self.get_command_handler(command) return handler(command) def get_command_handler(self, command): if command.command is None: return self.command_unhandled handler_name = "command_" + command.command.name handler = getattr(self, handler_name, None) return handler or self.command_unhandled def command_unhandled(self, command): logger.info("unhandled command", extra={"command": command}) def command_IP(self, command): logger.debug("interrupt process") # Currently we only ever read peer input during prompting, so # there's no process to interrupt. For bonus complexity there's a # fair chance the peer will send a timing mark request and discard # everything it gets before we acknowledge, so there's no point # sending a prompt unless we're sure they're ready. For no just # clear the input. self.line_buffer.clear() self.prompt_mgr.mark_interrupt() TimingMark = collections.namedtuple("TimingMark", ["option"])