Repaired websockets.
This commit is contained in:
		
							parent
							
								
									1de2c55966
								
							
						
					
					
						commit
						529606955a
					
				| @ -180,8 +180,8 @@ class Application(BaseApplication): | ||||
| 
 | ||||
| executor = ThreadPoolExecutor(max_workers=100) | ||||
| 
 | ||||
| loop = asyncio.get_event_loop() | ||||
| loop.set_default_executor(executor) | ||||
| #loop = asyncio.get_event_loop() | ||||
| #loop.set_default_executor(executor) | ||||
| 
 | ||||
| app = Application(db_path="sqlite:///snek.db") | ||||
| 
 | ||||
|  | ||||
| @ -13,80 +13,42 @@ commands = { | ||||
| } | ||||
| 
 | ||||
| class TerminalSession: | ||||
|      | ||||
|     async def ensure_process(self): | ||||
|         if self.process: | ||||
|             return | ||||
|         self.process = await asyncio.create_subprocess_exec( | ||||
|             *self.command.split(" "), | ||||
|             stdin=self.slave, | ||||
|             stdout=self.slave, | ||||
|             stderr=self.slave,bufsize=0, | ||||
|             universal_newlines=True | ||||
|         ) | ||||
| 
 | ||||
| 
 | ||||
|     def __init__(self,command): | ||||
|         self.master, self.slave = pty.openpty() | ||||
|         self.sockets =[] | ||||
|         self.buffer = b'' | ||||
|         self.process = None  | ||||
| 
 | ||||
|         #self.process = subprocess.Popen( | ||||
|         #    command.split(" "), | ||||
|         #    stdin=self.slave, | ||||
|         #    stdout=self.slave, | ||||
|         #    stderr=self.slave, | ||||
|         #    bufsize=0, | ||||
|         #    universal_newlines=True | ||||
|         #) | ||||
|         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 | ||||
|         ) | ||||
| 
 | ||||
|     async def add_websocket(self, ws): | ||||
|         asyncio.create_task(self.read_output(ws)) | ||||
| 
 | ||||
|     async def read_output(self, ws): | ||||
|         await self.ensure_process() | ||||
|         self.sockets.append(ws) | ||||
|         if len(self.sockets) > 1: | ||||
|             start = self.buffer.index(b'\n') | ||||
|             await ws.send_bytes(self.buffer[start:]) | ||||
|             return  | ||||
|         while True: | ||||
|         if len(self.sockets) > 1 and self.buffer: | ||||
|             start = 0 | ||||
|             try: | ||||
|                 async for data in self.process.stdout: | ||||
|                     if not data: | ||||
|                         break | ||||
|                     self.buffer += data | ||||
|                     if len(self.buffer) > 10000: | ||||
|                         self.buffer = self.buffer[:-10000] | ||||
|                     try: | ||||
|                         for ws in self.sockets: await ws.send_bytes(data)  # Send raw bytes for ANSI support | ||||
|                     except: | ||||
|                         self.sockets.remove(ws) | ||||
|             except: | ||||
|                 print("Terminating process") | ||||
|                 self.process.terminate() | ||||
|                 print("Terminated process") | ||||
|                 for ws in self.sockets: | ||||
|                     try: | ||||
|                         await ws.close() | ||||
|                     except Exception: | ||||
|                 start = self.history.index(b'\n') | ||||
|             except ValueError: | ||||
|                 pass  | ||||
|                 break | ||||
| 
 | ||||
|     async def read_outputa(self, ws): | ||||
|         loop = asyncio.get_event_loop() | ||||
|         self.sockets.append(ws) | ||||
|         if len(self.sockets) > 1: | ||||
|             start = self.buffer.index(b'\n') | ||||
|             await ws.send_bytes(self.buffer[start:]) | ||||
|             await ws.send_bytes(self.history[start:]) | ||||
|             return  | ||||
|         loop = asyncio.get_event_loop() | ||||
|         while True: | ||||
|             try: | ||||
|                 data = await loop.run_in_executor(None, os.read, self.master, 1024) | ||||
|                 if not data: | ||||
|                     break | ||||
|                 self.buffer += data | ||||
|                 if len(self.buffer) > 10000: | ||||
|                     self.buffer = self.buffer[:-10000] | ||||
|                 self.history += data | ||||
|                 if len(self.history) > self.history_size: | ||||
|                     self.history = self.history[:0-self.history_size] | ||||
|                 try: | ||||
|                     for ws in self.sockets: await ws.send_bytes(data)  # Send raw bytes for ANSI support | ||||
|                 except: | ||||
| @ -103,7 +65,10 @@ class TerminalSession: | ||||
|                 break | ||||
| 
 | ||||
|     async def write_input(self, data): | ||||
|         await self.ensure_process() | ||||
|         os.write(self.master, data.encode()) | ||||
| 
 | ||||
| 
 | ||||
|         try: | ||||
|             data = data.encode() | ||||
|         except AttributeError: | ||||
|             pass | ||||
|         await asyncio.get_event_loop().run_in_executor( | ||||
|             None, os.write, self.master, data | ||||
|         ) | ||||
|  | ||||
| @ -37,7 +37,8 @@ class TerminalSocketView(BaseView): | ||||
|         if not session: | ||||
|             self.user_sessions[user["uid"]] = TerminalSession(command=command) | ||||
|         session = self.user_sessions[user["uid"]]    | ||||
|         asyncio.create_task(session.read_output(ws))   | ||||
|         await session.add_websocket(ws) | ||||
|         #asyncio.create_task(session.read_output(ws))   | ||||
| 
 | ||||
|         async for msg in ws: | ||||
|             if msg.type == aiohttp.WSMsgType.BINARY: | ||||
|  | ||||
		Loading…
	
		Reference in New Issue
	
	Block a user