-
Notifications
You must be signed in to change notification settings - Fork 94
/
websocket.py
136 lines (109 loc) · 4.49 KB
/
websocket.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
"""Tornado websocket handler to serve a terminal interface.
"""
# Copyright (c) Jupyter Development Team
# Copyright (c) 2014, Ramalingam Saravanan <sarava@sarava.net>
# Distributed under the terms of the Simplified BSD License.
from __future__ import absolute_import, print_function
# Python3-friendly imports
import os
try:
from urllib.parse import urlparse
except ImportError:
from urlparse import urlparse
import json
import logging
import tornado.websocket
def _cast_unicode(s):
if isinstance(s, bytes):
return s.decode('utf-8')
return s
class TermSocket(tornado.websocket.WebSocketHandler):
"""Handler for a terminal websocket"""
def initialize(self, term_manager):
self.term_manager = term_manager
self.term_name = ""
self.size = (None, None)
self.terminal = None
self._logger = logging.getLogger(__name__)
self._user_command = ''
# Enable if the environment variable LOG_TERMINAL_OUTPUT is "true"
self._enable_output_logging = (str.lower(os.getenv("LOG_TERMINAL_OUTPUT", "false")) == "true")
def origin_check(self, origin=None):
"""Deprecated: backward-compat for terminado <= 0.5."""
return self.check_origin(origin or self.request.headers.get('Origin'))
def open(self, url_component=None):
"""Websocket connection opened.
Call our terminal manager to get a terminal, and connect to it as a
client.
"""
# Jupyter has a mixin to ping websockets and keep connections through
# proxies alive. Call super() to allow that to set up:
super(TermSocket, self).open(url_component)
self._logger.info("TermSocket.open: %s", url_component)
url_component = _cast_unicode(url_component)
self.term_name = url_component or 'tty'
self.terminal = self.term_manager.get_terminal(url_component)
self.terminal.clients.append(self)
self.send_json_message(["setup", {}])
self._logger.info("TermSocket.open: Opened %s", self.term_name)
# Now drain the preopen buffer, if reconnect.
buffered = ""
preopen_buffer = self.terminal.read_buffer.copy()
while True:
if not preopen_buffer:
break
s = preopen_buffer.popleft()
buffered += s
if buffered:
self.on_pty_read(buffered)
def on_pty_read(self, text):
"""Data read from pty; send to frontend"""
self.send_json_message(['stdout', text])
def send_json_message(self, content):
json_msg = json.dumps(content)
self.write_message(json_msg)
if self._enable_output_logging:
if content[0] == "stdout" and isinstance(content[1], str):
self.log_terminal_output(f'STDOUT: {content[1]}')
def on_message(self, message):
"""Handle incoming websocket message
We send JSON arrays, where the first element is a string indicating
what kind of message this is. Data associated with the message follows.
"""
##logging.info("TermSocket.on_message: %s - (%s) %s", self.term_name, type(message), len(message) if isinstance(message, bytes) else message[:250])
command = json.loads(message)
msg_type = command[0]
if msg_type == "stdin":
self.terminal.ptyproc.write(command[1])
if self._enable_output_logging:
if command[1] == '\r':
self.log_terminal_output(f'STDIN: {self._user_command}')
self._user_command = ''
else:
self._user_command += command[1]
elif msg_type == "set_size":
self.size = command[1:3]
self.terminal.resize_to_smallest()
def on_close(self):
"""Handle websocket closing.
Disconnect from our terminal, and tell the terminal manager we're
disconnecting.
"""
self._logger.info("Websocket closed")
if self.terminal:
self.terminal.clients.remove(self)
self.terminal.resize_to_smallest()
self.term_manager.client_disconnected(self)
def on_pty_died(self):
"""Terminal closed: tell the frontend, and close the socket.
"""
self.send_json_message(['disconnect', 1])
self.close()
self.terminal = None
def log_terminal_output(self, log: str = ''):
"""
Logs the terminal input/output
:param log: log line to write
:return:
"""
self._logger.debug(log)