Files
webssh/webssh/worker.py
T
Jocay c013a389fe feat: Implement directory validation and shell command quoting
- Added is_valid_directory function to validate directory paths, ensuring they do not contain control characters and are within a specified length.
- Introduced quote_shell_arg to safely quote shell arguments, preventing command injection.
- Created build_cd_command to generate a command for changing directories in a shell.
- Enhanced the LoginHandler to utilize a login rate limiter, preventing brute-force attacks by tracking failed login attempts.
- Implemented an EncodingCache to optimize encoding detection for SSH connections.
- Updated the UI to include an input field for specifying an initial directory upon login, with appropriate validation and hints.
- Added a quickbar in the terminal interface for easy access to copy and paste functionality.
- Introduced a toast notification system to provide feedback on copy actions.
- Refactored connection storage to encrypt passwords at rest, improving security.
- Updated various templates and styles to accommodate new features and improve user experience.
2026-08-10 00:07:52 +08:00

136 lines
3.9 KiB
Python

import logging
try:
import secrets
except ImportError:
secrets = None
import tornado.websocket
from uuid import uuid4
from tornado.ioloop import IOLoop
from tornado.iostream import _ERRNO_CONNRESET
from tornado.util import errno_from_exception
BUF_SIZE = 32 * 1024
clients = {} # {client_key: {id: worker}}
def clear_worker(worker, clients):
key = worker.client_key
workers = clients.get(key)
if not workers or worker.id not in workers:
return
workers.pop(worker.id)
if not workers:
clients.pop(key, None)
def recycle_worker(worker):
if worker.handler:
return
logging.warning('Recycling worker {}'.format(worker.id))
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 = self.gen_id()
self.data_to_dst = []
self.handler = None
self.mode = IOLoop.READ
self.closed = False
self.src_addr = None
self.client_key = None
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')
@classmethod
def gen_id(cls):
return secrets.token_urlsafe(nbytes=32) if secrets else uuid4().hex
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
if mode == IOLoop.WRITE:
self.loop.call_later(0.1, self, self.fd, IOLoop.WRITE)
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 self.chan.closed or 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, binary=True)
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 self.chan.closed or 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):
if self.closed:
return
self.closed = True
logging.info(
'Closing worker {} with reason: {}'.format(self.id, reason)
)
if self.handler:
self.loop.remove_handler(self.fd)
self.handler.close(reason=reason)
self.chan.close()
self.ssh.close()
logging.info('Connection to {}:{} lost'.format(*self.dst_addr))
clear_worker(self, clients)
logging.debug(clients)