386 lines
14 KiB
Python
386 lines
14 KiB
Python
import argparse
|
|
import asyncio
|
|
import contextlib
|
|
import errno
|
|
import fcntl
|
|
import os
|
|
import re
|
|
import signal
|
|
import sys
|
|
import termios
|
|
import threading
|
|
import tty
|
|
import queue
|
|
from types import TracebackType
|
|
from typing import Any, BinaryIO, Iterable, List, Optional, Tuple, Type
|
|
from typing_extensions import Literal
|
|
|
|
import serial
|
|
import serial_asyncio
|
|
|
|
|
|
MODE_RE = re.compile(r"^([856])([NEOMS])(1|1.5|2)$")
|
|
MODE_LOOKUP = {
|
|
"bytesize": {
|
|
"8": serial.EIGHTBITS, "5": serial.FIVEBITS, "6": serial.SIXBITS
|
|
},
|
|
"parity": {
|
|
"N": serial.PARITY_NONE, "E": serial.PARITY_EVEN, "O": serial.PARITY_ODD,
|
|
"M": serial.PARITY_MARK, "S": serial.PARITY_SPACE
|
|
},
|
|
"stopbits": {
|
|
"1": serial.STOPBITS_ONE, "1.5": serial.STOPBITS_ONE_POINT_FIVE, "2": serial.STOPBITS_TWO
|
|
}
|
|
}
|
|
|
|
|
|
class TtyRaw:
|
|
__slots__ = ["isatty", "infd", "outfd", "settings", "read_thread", "write_thread"]
|
|
isatty: bool
|
|
infd: int
|
|
outfd: int
|
|
settings: List[Any]
|
|
read_thread: Optional[threading.Thread]
|
|
write_thread: Optional[threading.Thread]
|
|
|
|
def __init__(self) -> None:
|
|
self.isatty = False
|
|
self.infd = 0
|
|
self.outfd = 0
|
|
self.settings = []
|
|
self.read_thread = None
|
|
self.write_thread = None
|
|
|
|
def __enter__(self) -> 'TtyRaw':
|
|
if sys.stdin.isatty():
|
|
self.isatty = True
|
|
self.infd = sys.stdin.fileno()
|
|
self.outfd = sys.stdout.fileno()
|
|
self.settings = termios.tcgetattr(self.infd)
|
|
tty.setraw(self.infd)
|
|
return self
|
|
|
|
def __exit__(self, exc_type: Optional[Type[BaseException]],
|
|
exc_value: Optional[BaseException],
|
|
exc_traceback: Optional[TracebackType]) -> Literal[False]:
|
|
if self.isatty:
|
|
termios.tcsetattr(self.infd, termios.TCSADRAIN, self.settings)
|
|
|
|
# unset nonblocking modes
|
|
flags = fcntl.fcntl(self.infd, fcntl.F_GETFL)
|
|
flags = flags & (~os.O_NONBLOCK)
|
|
fcntl.fcntl(self.infd, fcntl.F_SETFL, flags)
|
|
flags = fcntl.fcntl(self.outfd, fcntl.F_GETFL)
|
|
flags = flags & (~os.O_NONBLOCK)
|
|
fcntl.fcntl(self.outfd, fcntl.F_SETFL, flags)
|
|
return False
|
|
|
|
async def setup_async(self) -> Tuple[asyncio.StreamReader, asyncio.StreamWriter]:
|
|
# this is some UNIX bullshit honestly
|
|
# if a stdio stream is a tty we assume we have full control over it and thus can set it
|
|
# into nonblocking mode (optimal) and do real asyncio
|
|
# if it's _not_ a tty, it could be a pipe, and in the UNIX piping scheme, programs really
|
|
# don't like it when you set a pipe to nonblocking and they're not expecting that
|
|
# therefore in that case we have this really ugly workaround that offloads stdio to thread
|
|
# executors which sucks a lot. but, it means megacom now works as expected with regular
|
|
# UNIX shell commands
|
|
# annoyingly, shell piping is when you'd actually want the higher performance allowed by
|
|
# using asyncio, so if you have some very spammy output you'd like to capture the best thing
|
|
# is to keep it on the tty and use -l to also save it to a log file
|
|
|
|
loop = asyncio.get_event_loop()
|
|
if sys.stdin.isatty():
|
|
reader = asyncio.StreamReader()
|
|
reader_protocol = asyncio.StreamReaderProtocol(reader)
|
|
await loop.connect_read_pipe(lambda: reader_protocol, sys.stdin.buffer)
|
|
else:
|
|
reader = asyncio.StreamReader()
|
|
def reader_thread():
|
|
while True:
|
|
# XXX: if this is higher than 1, python treats it like it's a
|
|
# "read up to N bytes" thing and will buffer input indefinitely. idk how to
|
|
# turn that behavior off
|
|
buf = sys.stdin.buffer.read(1)
|
|
if len(buf) == 0:
|
|
break
|
|
loop.call_soon_threadsafe(lambda buf: reader.feed_data(buf), buf)
|
|
loop.call_soon_threadsafe(lambda: reader.feed_eof())
|
|
self.read_thread = threading.Thread(target=reader_thread, daemon=True)
|
|
self.read_thread.start()
|
|
|
|
if sys.stdout.isatty():
|
|
writer_transport, writer_protocol = await loop.connect_write_pipe(
|
|
asyncio.streams.FlowControlMixin, sys.stdout.buffer)
|
|
writer = asyncio.StreamWriter(writer_transport, writer_protocol, None, loop)
|
|
else:
|
|
class ThreadWriter(asyncio.StreamWriter):
|
|
__slots__ = ["transport", "queue", "_is_closing", "close_evt"]
|
|
transport: asyncio.BaseTransport
|
|
queue: queue.Queue[Optional[bytes]]
|
|
_is_closing: bool
|
|
close_evt: threading.Event
|
|
|
|
def __init__(self):
|
|
self.transport = None
|
|
self.queue = queue.Queue()
|
|
self._is_closing = False
|
|
self.close_evt = threading.Event()
|
|
|
|
def write(self, data: bytes) -> None:
|
|
self.queue.put(data)
|
|
|
|
def writelines(self, data: Iterable[bytes]) -> None:
|
|
for d in data:
|
|
self.queue.put(d)
|
|
|
|
def close(self) -> None:
|
|
if not self._is_closing:
|
|
self._is_closing = True
|
|
self.queue.put(None)
|
|
|
|
def can_write_eof(self) -> bool:
|
|
return True
|
|
|
|
def write_eof(self) -> None:
|
|
self.close()
|
|
|
|
def get_extra_info(self, name: str, default: Any = None) -> Any:
|
|
return default
|
|
|
|
async def drain(self) -> None:
|
|
await asyncio.to_thread(lambda: self.queue.join())
|
|
await asyncio.to_thread(lambda: sys.stdout.buffer.flush())
|
|
|
|
def is_closing(self) -> bool:
|
|
return self._is_closing
|
|
|
|
async def wait_closed(self) -> None:
|
|
await asyncio.to_thread(lambda: self.close_evt.wait())
|
|
|
|
def _writing_thread(self):
|
|
while True:
|
|
data = self.queue.get()
|
|
if data is None:
|
|
break
|
|
sys.stdout.buffer.write(data)
|
|
self.queue.task_done()
|
|
sys.stdout.buffer.flush()
|
|
self.close_evt.set()
|
|
|
|
writer = ThreadWriter()
|
|
self.write_thread = threading.Thread(target=writer._writing_thread, daemon=True)
|
|
self.write_thread.start()
|
|
|
|
return (reader, writer)
|
|
|
|
|
|
# CTRL-A
|
|
ESC_CHAR = b"\x01"
|
|
|
|
|
|
class KeycodeHandler:
|
|
__slots__ = ["exit_flag", "esc", "isatty"]
|
|
exit_flag: asyncio.Event
|
|
esc: bool
|
|
isatty: bool
|
|
|
|
def __init__(self, isatty: bool) -> None:
|
|
self.exit_flag = asyncio.Event()
|
|
self.esc = False
|
|
self.isatty = isatty
|
|
|
|
def process(self, byte: bytes) -> bytes:
|
|
# only translate or eat input if stdin is actually a tty
|
|
if not self.isatty:
|
|
return byte
|
|
|
|
if self.esc:
|
|
self.esc = False
|
|
if byte == b"q":
|
|
self.exit_flag.set()
|
|
return b""
|
|
elif byte == ESC_CHAR:
|
|
return ESC_CHAR
|
|
else:
|
|
return b""
|
|
|
|
if byte == ESC_CHAR:
|
|
self.esc = True
|
|
return b""
|
|
# i'm not super sure why this translation is necessary
|
|
# idk if there's more translations that are necessary too
|
|
elif byte == b"\x7f":
|
|
return b"\x08"
|
|
|
|
return byte
|
|
|
|
|
|
async def megacom(ttyraw: TtyRaw, tty: str, baud: int, mode: str, logfile: Optional[str]) -> None:
|
|
(stdin, stdout) = await ttyraw.setup_async()
|
|
|
|
m = MODE_RE.match(mode)
|
|
if m is None:
|
|
sys.stderr.write(f"invalid mode: {mode}\n")
|
|
sys.exit(1)
|
|
bytesize = MODE_LOOKUP["bytesize"][m.group(1)]
|
|
parity = MODE_LOOKUP["parity"][m.group(2)]
|
|
stopbits = MODE_LOOKUP["stopbits"][m.group(3)]
|
|
|
|
log: Optional[BinaryIO] = None
|
|
if logfile is not None:
|
|
try:
|
|
log = open(logfile, "wb")
|
|
except Exception as e:
|
|
sys.stderr.write(f"failed to open log file: {e}\n")
|
|
sys.exit(1)
|
|
|
|
loop = asyncio.get_event_loop()
|
|
keycodes = KeycodeHandler(ttyraw.isatty)
|
|
|
|
loop.add_signal_handler(signal.SIGINT, lambda: keycodes.exit_flag.set())
|
|
|
|
try:
|
|
return await megacom_main(stdin, stdout, tty, baud, bytesize, parity, stopbits, loop,
|
|
keycodes, log)
|
|
finally:
|
|
loop.remove_signal_handler(signal.SIGINT)
|
|
|
|
|
|
async def megacom_main(stdin: asyncio.StreamReader, stdout: asyncio.StreamWriter, tty: str,
|
|
baud: int, bytesize: Any, parity: Any, stopbits: Any,
|
|
loop: asyncio.AbstractEventLoop, keycodes: KeycodeHandler,
|
|
log: Optional[BinaryIO]) -> None:
|
|
printed_fnf = False
|
|
|
|
while True:
|
|
try:
|
|
(serialin, serialout) = await serial_asyncio.open_serial_connection(
|
|
loop=loop, url=tty, baudrate=baud, bytesize=bytesize,
|
|
parity=parity, stopbits=stopbits)
|
|
break
|
|
except serial.SerialException as e:
|
|
if e.errno == errno.ENOENT:
|
|
# the device could not be plugged in yet.. just wait
|
|
if not printed_fnf:
|
|
printed_fnf = True
|
|
stdout.write(f"waiting for {tty} to become available...\r\n".encode())
|
|
await stdout.drain()
|
|
else:
|
|
# permanent failure
|
|
stdout.write(f"failed to open port because: {e}\r\n".encode())
|
|
await stdout.drain()
|
|
return
|
|
except Exception as e:
|
|
# permanant failure
|
|
stdout.write(f"failed to open port because: {e}\r\n".encode())
|
|
await stdout.drain()
|
|
return
|
|
|
|
# wait a bit
|
|
start = loop.time()
|
|
while loop.time() - start < 2.0 and not keycodes.exit_flag.is_set():
|
|
timeout = loop.time() - start
|
|
try:
|
|
byte = await asyncio.wait_for(stdin.read(1), timeout=timeout)
|
|
keycodes.process(byte)
|
|
except asyncio.TimeoutError:
|
|
continue
|
|
|
|
if keycodes.exit_flag.is_set():
|
|
stdout.write(b"connection cancelled\r\n")
|
|
await stdout.drain()
|
|
return
|
|
|
|
def print_conn():
|
|
sys.stderr.write(f"megacom connected to {tty}\r\n")
|
|
sys.stderr.flush()
|
|
|
|
status_task = asyncio.to_thread(print_conn)
|
|
|
|
async def connect_pipe(pin: asyncio.StreamReader, pout: asyncio.StreamWriter,
|
|
ctrl: bool = False) -> None:
|
|
while not pin.at_eof():
|
|
c: bytes = await pin.read(1)
|
|
if len(c) == 0:
|
|
continue
|
|
|
|
if ctrl:
|
|
c = keycodes.process(c)
|
|
if len(c) == 0:
|
|
continue
|
|
else:
|
|
if log is not None:
|
|
log.write(c)
|
|
|
|
pout.write(c)
|
|
await pout.drain()
|
|
|
|
stdin_to_serial: asyncio.Task = asyncio.create_task(connect_pipe(stdin, serialout, True))
|
|
serial_to_stdout: asyncio.Task = asyncio.create_task(connect_pipe(serialin, stdout))
|
|
time_to_exit: asyncio.Task = asyncio.create_task(keycodes.exit_flag.wait())
|
|
|
|
await status_task
|
|
|
|
do_retry = False
|
|
|
|
def handle_done(task):
|
|
nonlocal do_retry
|
|
if task.done():
|
|
exc = task.exception()
|
|
if exc is not None:
|
|
stdout.write(f"\r\n\r\nmegacom encountered error: {exc}\r\n".encode())
|
|
if isinstance(exc, serial.SerialException):
|
|
do_retry = True
|
|
else:
|
|
task.result()
|
|
return task.done()
|
|
|
|
await asyncio.wait([time_to_exit, stdin_to_serial, serial_to_stdout],
|
|
return_when=asyncio.FIRST_COMPLETED)
|
|
if handle_done(time_to_exit):
|
|
pass
|
|
elif handle_done(stdin_to_serial):
|
|
pass
|
|
elif handle_done(serial_to_stdout):
|
|
pass
|
|
|
|
time_to_exit.cancel()
|
|
stdin_to_serial.cancel()
|
|
serial_to_stdout.cancel()
|
|
with contextlib.suppress(asyncio.CancelledError):
|
|
await time_to_exit
|
|
with contextlib.suppress(asyncio.CancelledError):
|
|
with contextlib.suppress(serial.SerialException):
|
|
await stdin_to_serial
|
|
with contextlib.suppress(asyncio.CancelledError):
|
|
with contextlib.suppress(serial.SerialException):
|
|
await serial_to_stdout
|
|
|
|
if not do_retry:
|
|
stdout.write(b"\r\n\r\nmegacom is exiting\r\n")
|
|
|
|
with contextlib.suppress(serial.SerialException):
|
|
await serialout.drain()
|
|
with contextlib.suppress(serial.SerialException):
|
|
serialout.close()
|
|
await stdout.drain()
|
|
|
|
if do_retry:
|
|
return await megacom_main(stdin, stdout, tty, baud, bytesize, parity, stopbits, loop,
|
|
keycodes, log)
|
|
|
|
|
|
def main() -> None:
|
|
parser = argparse.ArgumentParser(prog="megacom",
|
|
description="Alternative console-based UART client")
|
|
parser.add_argument("tty", type=str, default="/dev/ttyUSB0", nargs="?",
|
|
help="Path to UART device [/dev/ttyUSB0]")
|
|
parser.add_argument("-b", "--baud", type=int, default=115200, help="UART baud rate [115200]")
|
|
parser.add_argument("-m", "--mode", type=str, default="8N1", help="UART mode string [8N1]")
|
|
parser.add_argument("-l", "--logfile", type=str, default="", help="file to log to")
|
|
args = parser.parse_args()
|
|
|
|
with TtyRaw() as ttyraw:
|
|
asyncio.run(megacom(ttyraw, args.tty, args.baud, args.mode,
|
|
args.logfile if args.logfile != "" else None))
|