django_webshell/webshell/consumers.py
Morgan 'ARR\!' Allen acae46d111 init
2026-05-08 22:00:50 -07:00

122 lines
3.1 KiB
Python

import asyncio
import asyncssh
from asgiref.sync import async_to_sync
from channels.generic.websocket import AsyncJsonWebsocketConsumer
from channels.consumer import AsyncConsumer
from channels.layers import get_channel_layer
class WebShellSocketConsumer(AsyncJsonWebsocketConsumer):
def __init__(self):
super().__init__()
self.conn = None
async def connect(self):
super().connect()
if self.scope['session'].session_key is None:
await self.scope['session'].acreate()
await self.scope['session'].asave()
await self.accept()
print('accepted connection as', self.scope['session'].session_key)
async def receive_json(self, data):
data['owner'] = self.scope['session'].session_key
await self.channel_layer.group_add(data['owner'], self.channel_name)
await self.channel_layer.send('terminal', data)
async def notify(self, event):
await self.send_json({
'type': 'stdout',
'stdout': event['data']
})
class WebShellClient(asyncssh.SSHClient):
pass
class WebShellClientSession(asyncssh.SSHClientSession):
def data_received(self, data, datatype):
channel = get_channel_layer()
asyncio.ensure_future(channel.group_send(self.channel_name, {
'type': 'notify',
'data': data,
}))
def eof_received(self):
print('ssh connection done')
class WebShellWorker(AsyncConsumer):
def __init__(self):
self.sessions = {}
async def ssh(self, event):
host = event['host']
user = event['username']
pw = event['password']
if not host:
await self.send_json({
'error': 'no IP provided'
})
return
if not user:
await self.send_json({
'error': 'no username provided'
})
return
if not pw:
await self.send_json({
'error': 'no password provided'
})
return
# TODO
# Track these on Consumer
if not event['owner'] in self.sessions:
print('creating new session')
conn, client = await asyncssh.create_connection(
WebShellClient,
host,
username=user,
password=pw,
known_hosts=None,
options=asyncssh.SSHClientConnectionOptions(
#server_host_key_algs='ssh-rsa'
)
)
chan, session = await conn.create_session(
WebShellClientSession,
env={
'channel_name': event['owner'],
},
term_type='xterm'
)
session.channel_name = event['owner']
session.worker = self
self.sessions[event['owner']] = {
'connection': conn,
'client': client,
'channel': chan,
'session': session,
}
async def key(self, event):
self.sessions[event['owner']]['channel'].write(event['key'])