From 655e6a9c8a24359f0307771540528327d1688de2 Mon Sep 17 00:00:00 2001
From: haskal <haskal@awoo.systems>
Date: Wed, 12 May 2021 16:42:43 -0400
Subject: [PATCH] allow exiting during waiting phase

---
 megacom/__init__.py | 60 +++++++++++++++++++++++++++++++++++++++------
 1 file changed, 52 insertions(+), 8 deletions(-)

diff --git a/megacom/__init__.py b/megacom/__init__.py
index 4929823..9166ebb 100644
--- a/megacom/__init__.py
+++ b/megacom/__init__.py
@@ -68,6 +68,37 @@ async def setup_async() -> Tuple[asyncio.StreamReader, asyncio.StreamWriter]:
     return (reader, writer)
 
 
+# CTRL-A
+ESC_CHAR = b"\x01"
+
+
+class KeycodeHandler:
+    __slots__ = ["exit_flag", "esc"]
+    exit_flag: asyncio.Event
+    esc: bool
+
+    def __init__(self) -> None:
+        self.exit_flag = asyncio.Event()
+        self.esc = False
+
+    def process(self, byte: bytes) -> bytes:
+        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""
+
+        return byte
+
+
 async def megacom(tty: str, baud: int, mode: str) -> None:
     (stdin, stdout) = await setup_async()
 
@@ -78,12 +109,15 @@ async def megacom(tty: str, baud: int, mode: str) -> None:
     parity = MODE_LOOKUP["parity"][m.group(2)]
     stopbits = MODE_LOOKUP["stopbits"][m.group(3)]
 
+    loop = asyncio.get_event_loop()
+    keycodes = KeycodeHandler()
+
     printed_fnf = False
 
     while True:
         try:
             (serialin, serialout) = await serial_asyncio.open_serial_connection(
-                loop=asyncio.get_event_loop(), url=tty, baudrate=baud, bytesize=bytesize,
+                loop=loop, url=tty, baudrate=baud, bytesize=bytesize,
                 parity=parity, stopbits=stopbits)
             break
         except serial.SerialException as e:
@@ -103,17 +137,27 @@ async def megacom(tty: str, baud: int, mode: str) -> None:
             stdout.write(f"failed to open port because: {e}\r\n".encode())
             await stdout.drain()
             return
-        await asyncio.sleep(2)
+
+        # 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
 
     stdout.write(f"megacom connected to {tty}:{baud}:{mode}\r\n".encode())
     await stdout.drain()
 
-    exit_flag = asyncio.Event()
-
     async def connect_pipe(pin: asyncio.StreamReader, pout: asyncio.StreamWriter,
                            ctrl: bool = False) -> None:
-        nonlocal exit_flag
-
         esc: bool = False
         while not pin.at_eof():
             c: bytes = await pin.read(1)
@@ -126,7 +170,7 @@ async def megacom(tty: str, baud: int, mode: str) -> None:
                     pout.write(b"\x01")
                     await pout.drain()
                 elif c == b"q":
-                    exit_flag.set()
+                    keycodes.exit_flag.set()
                     break
 
             if c == b"\x01" and ctrl:
@@ -138,7 +182,7 @@ async def megacom(tty: str, baud: int, mode: str) -> None:
 
     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(exit_flag.wait())
+    time_to_exit: asyncio.Task = asyncio.create_task(keycodes.exit_flag.wait())
 
     do_retry = False