diff --git a/src/snek/system/terminal.py b/src/snek/system/terminal.py index 4d3781e..bd7e057 100644 --- a/src/snek/system/terminal.py +++ b/src/snek/system/terminal.py @@ -11,20 +11,40 @@ commands = { class TerminalSession: def __init__(self, command): - self.master, self.slave = pty.openpty() + self.master, self.slave = None,None + self.process = None self.sockets = [] self.history = b"" self.history_size = 1024 * 20 - self.process = subprocess.Popen( - command.split(" "), - stdin=self.slave, - stdout=self.slave, - stderr=self.slave, - bufsize=0, - universal_newlines=True, - ) + self.command = command + self.start_process(self.command) + + def start_process(self, command): + if not self.is_running(): + if self.master: + os.close(self.master) + os.close(self.slave) + self.master = None + self.slave = None + + self.master, self.slave = pty.openpty() + self.process = subprocess.Popen( + command.split(" "), + stdin=self.slave, + stdout=self.slave, + stderr=self.slave, + bufsize=0, + universal_newlines=True, + ) + + def is_running(self): + if not self.process: + return False + loop = asyncio.get_event_loop() + return self.process.poll() is None async def add_websocket(self, ws): + self.start_process(self.command) asyncio.create_task(self.read_output(ws)) async def read_output(self, ws): @@ -52,21 +72,37 @@ class TerminalSession: except: self.sockets.remove(ws) except Exception: - print("Terminating process") - self.process.terminate() - print("Terminated process") - for ws in self.sockets: - try: - await ws.close() - except Exception: - pass - break + await self.close() + break + + async def close(self): + print("Terminating process") + if self.process: + self.process.terminate() + self.process = None + if self.master: + os.close(self.master) + os.close(self.slave) + self.master = None + self.slave = None + + print("Terminated process") + for ws in self.sockets: + try: + await ws.close() + except Exception: + pass + self.sockets = [] async def write_input(self, data): try: data = data.encode() except AttributeError: pass - await asyncio.get_event_loop().run_in_executor( - None, os.write, self.master, data - ) + try: + await asyncio.get_event_loop().run_in_executor( + None, os.write, self.master, data + ) + except Exception as ex: + print(ex) + await self.close()