mirror of
https://github.com/huashengdun/webssh.git
synced 2026-08-20 12:46:39 +00:00
Prepare to write unit tests
This commit is contained in:
Sheng
@@ -0,0 +1,207 @@
|
||||
import io
|
||||
import logging
|
||||
import socket
|
||||
import threading
|
||||
import traceback
|
||||
import weakref
|
||||
import paramiko
|
||||
import tornado.web
|
||||
|
||||
from tornado.ioloop import IOLoop
|
||||
from worker import Worker, recycle_worker, workers
|
||||
|
||||
try:
|
||||
from concurrent.futures import Future
|
||||
except ImportError:
|
||||
from tornado.concurrent import Future
|
||||
|
||||
|
||||
DELAY = 3
|
||||
|
||||
|
||||
class MixinHandler(object):
|
||||
|
||||
def get_real_client_addr(self):
|
||||
ip = self.request.headers.get('X-Real-Ip')
|
||||
port = self.request.headers.get('X-Real-Port')
|
||||
|
||||
if ip is None and port is None:
|
||||
return
|
||||
|
||||
try:
|
||||
port = int(port)
|
||||
except (TypeError, ValueError):
|
||||
pass
|
||||
else:
|
||||
if ip: # does not validate ip and port here
|
||||
return (ip, port)
|
||||
logging.warn('Bad nginx configuration.')
|
||||
|
||||
|
||||
class IndexHandler(MixinHandler, tornado.web.RequestHandler):
|
||||
|
||||
def initialize(self, loop, policy, host_keys_settings):
|
||||
self.loop = loop
|
||||
self.policy = policy
|
||||
self.host_keys_settings = host_keys_settings
|
||||
|
||||
def get_privatekey(self):
|
||||
try:
|
||||
data = self.request.files.get('privatekey')[0]['body']
|
||||
except TypeError:
|
||||
return
|
||||
return data.decode('utf-8')
|
||||
|
||||
def get_specific_pkey(self, pkeycls, privatekey, password):
|
||||
logging.info('Trying {}'.format(pkeycls.__name__))
|
||||
try:
|
||||
pkey = pkeycls.from_private_key(io.StringIO(privatekey),
|
||||
password=password)
|
||||
except paramiko.PasswordRequiredException:
|
||||
raise ValueError('Need password to decrypt the private key.')
|
||||
except paramiko.SSHException:
|
||||
pass
|
||||
else:
|
||||
return pkey
|
||||
|
||||
def get_pkey_obj(self, privatekey, password):
|
||||
password = password.encode('utf-8') if password else None
|
||||
|
||||
pkey = self.get_specific_pkey(paramiko.RSAKey, privatekey, password)\
|
||||
or self.get_specific_pkey(paramiko.DSSKey, privatekey, password)\
|
||||
or self.get_specific_pkey(paramiko.ECDSAKey, privatekey, password)\
|
||||
or self.get_specific_pkey(paramiko.Ed25519Key, privatekey,
|
||||
password)
|
||||
if not pkey:
|
||||
raise ValueError('Not a valid private key file or '
|
||||
'wrong password for decrypting the private key.')
|
||||
return pkey
|
||||
|
||||
def get_port(self):
|
||||
value = self.get_value('port')
|
||||
try:
|
||||
port = int(value)
|
||||
except ValueError:
|
||||
port = 0
|
||||
|
||||
if 0 < port < 65536:
|
||||
return port
|
||||
|
||||
raise ValueError('Invalid port {}'.format(value))
|
||||
|
||||
def get_value(self, name):
|
||||
value = self.get_argument(name)
|
||||
if not value:
|
||||
raise ValueError('Empty {}'.format(name))
|
||||
return value
|
||||
|
||||
def get_args(self):
|
||||
hostname = self.get_value('hostname')
|
||||
port = self.get_port()
|
||||
username = self.get_value('username')
|
||||
password = self.get_argument('password')
|
||||
privatekey = self.get_privatekey()
|
||||
pkey = self.get_pkey_obj(privatekey, password) if privatekey else None
|
||||
args = (hostname, port, username, password, pkey)
|
||||
logging.debug(args)
|
||||
return args
|
||||
|
||||
def get_client_addr(self):
|
||||
return self.get_real_client_addr() or self.request.connection.stream.\
|
||||
socket.getpeername()
|
||||
|
||||
def ssh_connect(self):
|
||||
ssh = paramiko.SSHClient()
|
||||
ssh._system_host_keys = self.host_keys_settings['system_host_keys']
|
||||
ssh._host_keys = self.host_keys_settings['host_keys']
|
||||
ssh._host_keys_filename = self.host_keys_settings['host_keys_filename']
|
||||
ssh.set_missing_host_key_policy(self.policy)
|
||||
|
||||
args = self.get_args()
|
||||
dst_addr = (args[0], args[1])
|
||||
logging.info('Connecting to {}:{}'.format(*dst_addr))
|
||||
|
||||
try:
|
||||
ssh.connect(*args, timeout=6)
|
||||
except socket.error:
|
||||
raise ValueError('Unable to connect to {}:{}'.format(*dst_addr))
|
||||
except paramiko.BadAuthenticationType:
|
||||
raise ValueError('SSH authentication failed.')
|
||||
except paramiko.BadHostKeyException:
|
||||
raise ValueError('Bad host key.')
|
||||
|
||||
chan = ssh.invoke_shell(term='xterm')
|
||||
chan.setblocking(0)
|
||||
worker = Worker(self.loop, ssh, chan, dst_addr)
|
||||
worker.src_addr = self.get_client_addr()
|
||||
return worker
|
||||
|
||||
def ssh_connect_wrapped(self, future):
|
||||
try:
|
||||
worker = self.ssh_connect()
|
||||
except Exception as exc:
|
||||
logging.error(traceback.format_exc())
|
||||
future.set_exception(exc)
|
||||
else:
|
||||
future.set_result(worker)
|
||||
|
||||
def get(self):
|
||||
self.render('index.html')
|
||||
|
||||
@tornado.gen.coroutine
|
||||
def post(self):
|
||||
worker_id = None
|
||||
status = None
|
||||
|
||||
future = Future()
|
||||
t = threading.Thread(target=self.ssh_connect_wrapped, args=(future,))
|
||||
t.setDaemon(True)
|
||||
t.start()
|
||||
|
||||
try:
|
||||
worker = yield future
|
||||
except Exception as exc:
|
||||
status = str(exc)
|
||||
else:
|
||||
worker_id = worker.id
|
||||
workers[worker_id] = worker
|
||||
self.loop.call_later(DELAY, recycle_worker, worker)
|
||||
|
||||
self.write(dict(id=worker_id, status=status))
|
||||
|
||||
|
||||
class WsockHandler(MixinHandler, tornado.websocket.WebSocketHandler):
|
||||
|
||||
def initialize(self, loop):
|
||||
self.loop = loop
|
||||
self.worker_ref = None
|
||||
|
||||
def get_client_addr(self):
|
||||
return self.get_real_client_addr() or self.stream.socket.getpeername()
|
||||
|
||||
def open(self):
|
||||
self.src_addr = self.get_client_addr()
|
||||
logging.info('Connected from {}:{}'.format(*self.src_addr))
|
||||
worker = workers.get(self.get_argument('id'))
|
||||
if worker and worker.src_addr[0] == self.src_addr[0]:
|
||||
workers.pop(worker.id)
|
||||
self.set_nodelay(True)
|
||||
worker.set_handler(self)
|
||||
self.worker_ref = weakref.ref(worker)
|
||||
self.loop.add_handler(worker.fd, worker, IOLoop.READ)
|
||||
else:
|
||||
self.close(reason='Websocket authentication failed.')
|
||||
|
||||
def on_message(self, message):
|
||||
logging.debug('{!r} from {}:{}'.format(message, *self.src_addr))
|
||||
worker = self.worker_ref()
|
||||
worker.data_to_dst.append(message)
|
||||
worker.on_write()
|
||||
|
||||
def on_close(self):
|
||||
logging.info('Disconnected from {}:{}'.format(*self.src_addr))
|
||||
worker = self.worker_ref() if self.worker_ref else None
|
||||
if worker:
|
||||
if self.close_reason is None:
|
||||
self.close_reason = 'client disconnected'
|
||||
worker.close(reason=self.close_reason)
|
||||
@@ -0,0 +1,36 @@
|
||||
import logging
|
||||
import tornado.web
|
||||
import tornado.ioloop
|
||||
|
||||
from tornado.options import parse_command_line, options
|
||||
from handler import IndexHandler, WsockHandler
|
||||
from settings import (get_app_settings, get_host_keys_settings,
|
||||
get_policy_setting)
|
||||
|
||||
|
||||
def make_app(loop, policy, host_keys_settings, app_settings):
|
||||
handlers = [
|
||||
(r'/', IndexHandler, dict(loop=loop, policy=policy,
|
||||
host_keys_settings=host_keys_settings)),
|
||||
(r'/ws', WsockHandler, dict(loop=loop))
|
||||
]
|
||||
|
||||
app = tornado.web.Application(handlers, **app_settings)
|
||||
return app
|
||||
|
||||
|
||||
def main():
|
||||
parse_command_line()
|
||||
app_settings = get_app_settings(options)
|
||||
host_keys_settings = get_host_keys_settings(options)
|
||||
policy = get_policy_setting(options, host_keys_settings)
|
||||
|
||||
loop = tornado.ioloop.IOLoop.current()
|
||||
app = make_app(loop, policy, host_keys_settings, app_settings)
|
||||
app.listen(options.port, options.address)
|
||||
logging.info('Listening on {}:{}'.format(options.address, options.port))
|
||||
loop.start()
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
main()
|
||||
@@ -0,0 +1,83 @@
|
||||
import logging
|
||||
import os.path
|
||||
import threading
|
||||
import paramiko
|
||||
|
||||
|
||||
def load_host_keys(path):
|
||||
if os.path.exists(path) and os.path.isfile(path):
|
||||
return paramiko.hostkeys.HostKeys(filename=path)
|
||||
return paramiko.hostkeys.HostKeys()
|
||||
|
||||
|
||||
def get_policy_dictionary():
|
||||
dic = {
|
||||
k.lower(): v for k, v in vars(paramiko.client).items() if type(v)
|
||||
is type and issubclass(v, paramiko.client.MissingHostKeyPolicy)
|
||||
and v is not paramiko.client.MissingHostKeyPolicy
|
||||
}
|
||||
return dic
|
||||
|
||||
|
||||
def get_policy_class(policy):
|
||||
origin_policy = policy
|
||||
policy = policy.lower()
|
||||
if not policy.endswith('policy'):
|
||||
policy += 'policy'
|
||||
|
||||
dic = get_policy_dictionary()
|
||||
logging.debug(dic)
|
||||
|
||||
try:
|
||||
cls = dic[policy]
|
||||
except KeyError:
|
||||
raise ValueError('Unknown policy {!r}'.format(origin_policy))
|
||||
return cls
|
||||
|
||||
|
||||
def check_policy_setting(policy_class, host_keys_settings):
|
||||
host_keys = host_keys_settings['host_keys']
|
||||
host_keys_filename = host_keys_settings['host_keys_filename']
|
||||
system_host_keys = host_keys_settings['system_host_keys']
|
||||
|
||||
if policy_class is paramiko.client.AutoAddPolicy:
|
||||
host_keys.save(host_keys_filename) # for permission test
|
||||
elif policy_class is paramiko.client.RejectPolicy:
|
||||
if not host_keys and not system_host_keys:
|
||||
raise ValueError(
|
||||
'Reject policy could not be used without host keys.'
|
||||
)
|
||||
|
||||
|
||||
class AutoAddPolicy(paramiko.client.MissingHostKeyPolicy):
|
||||
"""
|
||||
thread-safe AutoAddPolicy
|
||||
"""
|
||||
lock = threading.Lock()
|
||||
|
||||
def is_missing_host_key(self, client, hostname, key):
|
||||
k = client._host_keys.lookup(hostname)
|
||||
if k is None:
|
||||
return True
|
||||
host_key = k.get(key.get_name(), None)
|
||||
if host_key is None:
|
||||
return True
|
||||
if host_key != key:
|
||||
raise paramiko.BadHostKeyException(hostname, key, host_key)
|
||||
|
||||
def missing_host_key(self, client, hostname, key):
|
||||
with self.lock:
|
||||
if self.is_missing_host_key(client, hostname, key):
|
||||
keytype = key.get_name()
|
||||
logging.info(
|
||||
'Adding {} host key for {}'.format(keytype, hostname)
|
||||
)
|
||||
client._host_keys._entries.append(
|
||||
paramiko.hostkeys.HostKeyEntry([hostname], key)
|
||||
)
|
||||
|
||||
with open(client._host_keys_filename, 'a') as f:
|
||||
f.write('{} {} {}\n'.format(
|
||||
hostname, keytype, key.get_base64()
|
||||
))
|
||||
paramiko.client.AutoAddPolicy = AutoAddPolicy
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 68 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 83 KiB |
@@ -0,0 +1,57 @@
|
||||
import logging
|
||||
import os.path
|
||||
import uuid
|
||||
|
||||
from tornado.options import define
|
||||
from policy import load_host_keys, get_policy_class, check_policy_setting
|
||||
|
||||
|
||||
define('address', default='127.0.0.1', help='listen address')
|
||||
define('port', default=8888, help='listen port', type=int)
|
||||
define('debug', default=False, help='debug mode', type=bool)
|
||||
define('policy', default='warning',
|
||||
help='missing host key policy, reject|autoadd|warning')
|
||||
define('hostFile', default='', help='User defined host keys file')
|
||||
define('sysHostFile', default='', help='System wide host keys file')
|
||||
|
||||
|
||||
base_dir = os.path.dirname(__file__)
|
||||
|
||||
|
||||
def get_app_settings(options):
|
||||
settings = dict(
|
||||
template_path=os.path.join(base_dir, 'templates'),
|
||||
static_path=os.path.join(base_dir, 'static'),
|
||||
cookie_secret=uuid.uuid4().hex,
|
||||
xsrf_cookies=True,
|
||||
debug=options.debug
|
||||
)
|
||||
return settings
|
||||
|
||||
|
||||
def get_host_keys_settings(options):
|
||||
if not options.hostFile:
|
||||
host_keys_filename = os.path.join(base_dir, 'known_hosts')
|
||||
else:
|
||||
host_keys_filename = options.hostFile
|
||||
host_keys = load_host_keys(host_keys_filename)
|
||||
|
||||
if not options.sysHostFile:
|
||||
filename = os.path.expanduser('~/.ssh/known_hosts')
|
||||
else:
|
||||
filename = options.sysHostFile
|
||||
system_host_keys = load_host_keys(filename)
|
||||
|
||||
settings = dict(
|
||||
host_keys=host_keys,
|
||||
system_host_keys=system_host_keys,
|
||||
host_keys_filename=host_keys_filename
|
||||
)
|
||||
return settings
|
||||
|
||||
|
||||
def get_policy_setting(options, host_keys_settings):
|
||||
policy_class = get_policy_class(options.policy)
|
||||
logging.info(policy_class.__name__)
|
||||
check_policy_setting(policy_class, host_keys_settings)
|
||||
return policy_class()
|
||||
Vendored
+7
File diff suppressed because one or more lines are too long
Vendored
+2
@@ -0,0 +1,2 @@
|
||||
.xterm.fullscreen{position:fixed;top:0;bottom:0;left:0;right:0;width:auto;height:auto;z-index:255}
|
||||
/*# sourceMappingURL=fullscreen.min.css.map */
|
||||
Vendored
+2
File diff suppressed because one or more lines are too long
Binary file not shown.
|
After Width: | Height: | Size: 5.8 KiB |
Vendored
+7
File diff suppressed because one or more lines are too long
Vendored
+1
@@ -0,0 +1 @@
|
||||
!function(e){"object"==typeof exports&&"object"==typeof module?module.exports=e(require("../../xterm")):"function"==typeof define?define(["../../xterm"],e):e(window.Terminal)}(function(e){var t={};return t.toggleFullScreen=function(e,t){var n;n=void 0===t?e.element.classList.contains("fullscreen")?"remove":"add":t?"add":"remove",e.element.classList[n]("fullscreen")},e.prototype.toggleFullscreen=function(e){t.toggleFullScreen(this,e)},t});
|
||||
Vendored
+2
File diff suppressed because one or more lines are too long
@@ -0,0 +1,119 @@
|
||||
jQuery(function($){
|
||||
|
||||
var status = $('#status'),
|
||||
btn = $('.btn-primary'),
|
||||
style = {};
|
||||
|
||||
$('form#connect').submit(function(event) {
|
||||
event.preventDefault();
|
||||
|
||||
var form = $(this),
|
||||
url = form.attr('action'),
|
||||
type = form.attr('type'),
|
||||
data = new FormData(this);
|
||||
|
||||
if (!data.get('hostname') || !data.get('port') || !data.get('username')) {
|
||||
status.text('Hostname, port and username are required.');
|
||||
return;
|
||||
}
|
||||
|
||||
var pk = data.get('privatekey');
|
||||
if (pk && pk.size > 16384) {
|
||||
status.text('Key size exceeds maximum value.');
|
||||
return;
|
||||
}
|
||||
|
||||
status.text('');
|
||||
btn.prop('disabled', true);
|
||||
|
||||
$.ajax({
|
||||
url: url,
|
||||
type: type,
|
||||
data: data,
|
||||
success: callback,
|
||||
cache: false,
|
||||
contentType: false,
|
||||
processData: false
|
||||
});
|
||||
|
||||
});
|
||||
|
||||
function parse_xterm_style() {
|
||||
var text = $('.xterm-helpers style').text();
|
||||
var arr = text.split('xterm-normal-char{width:');
|
||||
style.width = parseInt(arr[1]) + 1;
|
||||
arr = text.split('div{height:');
|
||||
style.height = parseInt(arr[1]);
|
||||
}
|
||||
|
||||
function current_geometry() {
|
||||
if (!style.width || !style.height) {
|
||||
parse_xterm_style();
|
||||
}
|
||||
cols = parseInt(window.innerWidth / style.width);
|
||||
rows = parseInt(window.innerHeight / style.height);
|
||||
return [cols, rows];
|
||||
}
|
||||
|
||||
|
||||
function callback(msg) {
|
||||
// console.log(msg);
|
||||
if (msg.status) {
|
||||
status.text(msg.status);
|
||||
setTimeout(function(){
|
||||
btn.prop('disabled', false);
|
||||
}, 3000);
|
||||
return;
|
||||
}
|
||||
|
||||
var ws_url = window.location.href.replace('http', 'ws'),
|
||||
join = (ws_url[ws_url.length-1] == '/' ? '' : '/'),
|
||||
url = ws_url + join + 'ws?id=' + msg.id,
|
||||
socket = new WebSocket(url),
|
||||
terminal = document.getElementById('#terminal'),
|
||||
geometry = current_geometry();
|
||||
term = new Terminal({
|
||||
cursorBlink: true,
|
||||
cols: geometry[0],
|
||||
rows: geometry[1]
|
||||
});
|
||||
|
||||
console.log(url);
|
||||
term.on('data', function(data) {
|
||||
// console.log(data);
|
||||
socket.send(data);
|
||||
});
|
||||
|
||||
socket.onopen = function(e) {
|
||||
$('.container').hide();
|
||||
term.open(terminal, true);
|
||||
term.toggleFullscreen(true);
|
||||
};
|
||||
|
||||
socket.onmessage = function(msg) {
|
||||
// console.log(msg);
|
||||
term.write(msg.data);
|
||||
};
|
||||
|
||||
socket.onerror = function(e) {
|
||||
console.log(e);
|
||||
};
|
||||
|
||||
socket.onclose = function(e) {
|
||||
console.log(e);
|
||||
term.destroy();
|
||||
$('.container').show();
|
||||
status.text(e.reason);
|
||||
btn.prop('disabled', false);
|
||||
};
|
||||
}
|
||||
|
||||
$(window).resize(function(){
|
||||
if (typeof term != 'undefined') {
|
||||
geometry = current_geometry();
|
||||
term.geometry = geometry;
|
||||
term.resize(geometry[0], geometry[1]);
|
||||
}
|
||||
});
|
||||
|
||||
});
|
||||
Vendored
+5
File diff suppressed because one or more lines are too long
Vendored
+1
File diff suppressed because one or more lines are too long
@@ -0,0 +1,72 @@
|
||||
<!DOCTYPE html>
|
||||
<html lang="en">
|
||||
<head>
|
||||
<meta charset="UTF-8">
|
||||
<title> WebSSH </title>
|
||||
<link href="static/img/favicon.png" rel="icon" type="image/png">
|
||||
<link href="static/css/bootstrap.min.css" rel="stylesheet" type="text/css"/>
|
||||
<link href="static/css/xterm.min.css" rel="stylesheet" type="text/css"/>
|
||||
<link href="static/css/fullscreen.min.css" rel="stylesheet" type="text/css"/>
|
||||
<style>
|
||||
.row {
|
||||
margin-top: 20px;
|
||||
margin-bottom: 10px;
|
||||
}
|
||||
|
||||
.container {
|
||||
margin-top: 20px;
|
||||
}
|
||||
|
||||
</style>
|
||||
</head>
|
||||
<body>
|
||||
<div class="container">
|
||||
<form id="connect" action="" type="post" enctype="multipart/form-data">
|
||||
<div class="row">
|
||||
<div class="col">
|
||||
<label for="Hostname">Hostname</label>
|
||||
<input class="form-control" type="text" name="hostname" value="">
|
||||
</div>
|
||||
<div class="col">
|
||||
<label for="Port">Port</label>
|
||||
<input class="form-control" type="text" name="port" value="">
|
||||
</div>
|
||||
</div>
|
||||
<div class="row">
|
||||
<div class="col">
|
||||
<label for="Username">Username</label>
|
||||
<input class="form-control" type="text" name="username" value="">
|
||||
</div>
|
||||
<div class="col">
|
||||
<label for="Username">Private Key</label>
|
||||
<input class="form-control" type="file" name="privatekey" value="">
|
||||
</div>
|
||||
</div>
|
||||
<div class="row">
|
||||
<div class="col">
|
||||
<label for="Password">Password</label>
|
||||
<input class="form-control" type="password" name="password" placeholder="" value="">
|
||||
</div>
|
||||
<div class="col">
|
||||
If Private Key is chosen, password will be used to decrypt the Private Key if it is encrypted, otherwise used as the password of username.
|
||||
</div>
|
||||
</div>
|
||||
{% module xsrf_form_html() %}
|
||||
<button type="submit" class="btn btn-primary">Connect</button>
|
||||
</form>
|
||||
|
||||
</div>
|
||||
|
||||
<div class="container">
|
||||
<div id="status" style="color: red;"></div>
|
||||
<div id="terminal"></div>
|
||||
</div>
|
||||
|
||||
<script src="static/js/jquery.min.js"></script>
|
||||
<script src="static/js/popper.min.js"></script>
|
||||
<script src="static/js/bootstrap.min.js"></script>
|
||||
<script src="static/js/xterm.min.js"></script>
|
||||
<script src="static/js/fullscreen.min.js"></script>
|
||||
<script src="static/js/main.js"></script>
|
||||
</body>
|
||||
</html>
|
||||
@@ -0,0 +1,104 @@
|
||||
import logging
|
||||
import tornado.websocket
|
||||
|
||||
from tornado.ioloop import IOLoop
|
||||
from tornado.iostream import _ERRNO_CONNRESET
|
||||
from tornado.util import errno_from_exception
|
||||
|
||||
|
||||
BUF_SIZE = 1024
|
||||
workers = {}
|
||||
|
||||
|
||||
def recycle_worker(worker):
|
||||
if worker.handler:
|
||||
return
|
||||
logging.warn('Recycling worker {}'.format(worker.id))
|
||||
workers.pop(worker.id, None)
|
||||
worker.close(reason='worker recycled')
|
||||
|
||||
|
||||
class Worker(object):
|
||||
def __init__(self, loop, ssh, chan, dst_addr):
|
||||
self.loop = loop
|
||||
self.ssh = ssh
|
||||
self.chan = chan
|
||||
self.dst_addr = dst_addr
|
||||
self.fd = chan.fileno()
|
||||
self.id = str(id(self))
|
||||
self.data_to_dst = []
|
||||
self.handler = None
|
||||
self.mode = IOLoop.READ
|
||||
|
||||
def __call__(self, fd, events):
|
||||
if events & IOLoop.READ:
|
||||
self.on_read()
|
||||
if events & IOLoop.WRITE:
|
||||
self.on_write()
|
||||
if events & IOLoop.ERROR:
|
||||
self.close(reason='error event occurred')
|
||||
|
||||
def set_handler(self, handler):
|
||||
if not self.handler:
|
||||
self.handler = handler
|
||||
|
||||
def update_handler(self, mode):
|
||||
if self.mode != mode:
|
||||
self.loop.update_handler(self.fd, mode)
|
||||
self.mode = mode
|
||||
|
||||
def on_read(self):
|
||||
logging.debug('worker {} on read'.format(self.id))
|
||||
try:
|
||||
data = self.chan.recv(BUF_SIZE)
|
||||
except (OSError, IOError) as e:
|
||||
logging.error(e)
|
||||
if errno_from_exception(e) in _ERRNO_CONNRESET:
|
||||
self.close(reason='chan error on reading')
|
||||
else:
|
||||
logging.debug('{!r} from {}:{}'.format(data, *self.dst_addr))
|
||||
if not data:
|
||||
self.close(reason='chan closed')
|
||||
return
|
||||
|
||||
logging.debug('{!r} to {}:{}'.format(data, *self.handler.src_addr))
|
||||
try:
|
||||
self.handler.write_message(data)
|
||||
except tornado.websocket.WebSocketClosedError:
|
||||
self.close(reason='websocket closed')
|
||||
|
||||
def on_write(self):
|
||||
logging.debug('worker {} on write'.format(self.id))
|
||||
if not self.data_to_dst:
|
||||
return
|
||||
|
||||
data = ''.join(self.data_to_dst)
|
||||
logging.debug('{!r} to {}:{}'.format(data, *self.dst_addr))
|
||||
|
||||
try:
|
||||
sent = self.chan.send(data)
|
||||
except (OSError, IOError) as e:
|
||||
logging.error(e)
|
||||
if errno_from_exception(e) in _ERRNO_CONNRESET:
|
||||
self.close(reason='chan error on writing')
|
||||
else:
|
||||
self.update_handler(IOLoop.WRITE)
|
||||
else:
|
||||
self.data_to_dst = []
|
||||
data = data[sent:]
|
||||
if data:
|
||||
self.data_to_dst.append(data)
|
||||
self.update_handler(IOLoop.WRITE)
|
||||
else:
|
||||
self.update_handler(IOLoop.READ)
|
||||
|
||||
def close(self, reason=None):
|
||||
logging.info(
|
||||
'Closing worker {} with reason {}'.format(self.id, reason)
|
||||
)
|
||||
if self.handler:
|
||||
self.loop.remove_handler(self.fd)
|
||||
self.handler.close()
|
||||
self.chan.close()
|
||||
self.ssh.close()
|
||||
logging.info('Connection to {}:{} lost'.format(*self.dst_addr))
|
||||
Reference in New Issue
Block a user