diff --git a/spug_api/consumer/consumers.py b/spug_api/consumer/consumers.py index 3c92d4e..34715df 100644 --- a/spug_api/consumer/consumers.py +++ b/spug_api/consumer/consumers.py @@ -1,43 +1,18 @@ # Copyright: (c) OpenSpug Organization. https://github.com/openspug/spug # Copyright: (c) # Released under the AGPL-3.0 License. -from channels.generic.websocket import WebsocketConsumer from django.conf import settings from django_redis import get_redis_connection from asgiref.sync import async_to_sync from apps.host.models import Host +from consumer.utils import BaseConsumer from apps.account.utils import has_host_perm from threading import Thread import time import json -class ExecConsumer(WebsocketConsumer): - def __init__(self, *args, **kwargs): - super().__init__(*args, **kwargs) - self.token = self.scope['url_route']['kwargs']['token'] - self.rds = get_redis_connection() - - def connect(self): - self.accept() - - def disconnect(self, code): - self.rds.close() - - def get_response(self): - response = self.rds.brpop(self.token, timeout=5) - return response[1] if response else None - - def receive(self, **kwargs): - response = self.get_response() - while response: - data = response.decode() - self.send(text_data=data) - response = self.get_response() - self.send(text_data='pong') - - -class ComConsumer(WebsocketConsumer): +class ComConsumer(BaseConsumer): def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) token = self.scope['url_route']['kwargs']['token'] @@ -52,9 +27,6 @@ class ComConsumer(WebsocketConsumer): raise TypeError(f'unknown module for {module}') self.rds = get_redis_connection() - def connect(self): - self.accept() - def disconnect(self, code): self.rds.close() @@ -78,18 +50,17 @@ class ComConsumer(WebsocketConsumer): self.send(text_data='pong') -class SSHConsumer(WebsocketConsumer): +class SSHConsumer(BaseConsumer): def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) - self.user = self.scope['user'] self.id = self.scope['url_route']['kwargs']['id'] self.chan = None self.ssh = None def loop_read(self): + is_ready = False while True: data = self.chan.recv(32 * 1024) - # print('read: {!r}'.format(data)) if not data: self.close(3333) break @@ -97,6 +68,9 @@ class SSHConsumer(WebsocketConsumer): text = data.decode() except UnicodeDecodeError: text = data.decode(encoding='GBK', errors='ignore') + if not is_ready: + self.send(text_data='\033[2J\033[3J\033[1;1H') + is_ready = True self.send(text_data=text) def receive(self, text_data=None, bytes_data=None): @@ -116,34 +90,28 @@ class SSHConsumer(WebsocketConsumer): if self.ssh: self.ssh.close() - def connect(self): + def init(self): if has_host_perm(self.user, self.id): - self.accept() - self._init() + self.send(text_data='\r\n正在连接至主机 ...') + host = Host.objects.filter(pk=self.id).first() + if not host: + return self.close_with_message('未找到指定主机,请刷新页面重试。') + + try: + self.ssh = host.get_ssh().get_client() + except Exception as e: + return self.close_with_message(f'连接主机失败: {e}') + + self.chan = self.ssh.invoke_shell(term='xterm') + self.chan.transport.set_keepalive(30) + Thread(target=self.loop_read).start() else: - self.close() - - def _init(self): - self.send(text_data='\r\33[KConnecting ...\r') - host = Host.objects.filter(pk=self.id).first() - if not host: - self.send(text_data='Unknown host\r\n') - self.close() - try: - self.ssh = host.get_ssh().get_client() - except Exception as e: - self.send(text_data=f'Exception: {e}\r\n'.encode()) - self.close() - return - self.chan = self.ssh.invoke_shell(term='xterm') - self.chan.transport.set_keepalive(30) - Thread(target=self.loop_read).start() + self.close_with_message('你当前无权限操作该主机,请联系管理员授权。') -class NotifyConsumer(WebsocketConsumer): - def connect(self): +class NotifyConsumer(BaseConsumer): + def init(self): async_to_sync(self.channel_layer.group_add)('notify', self.channel_name) - self.accept() def disconnect(self, code): async_to_sync(self.channel_layer.group_discard)('notify', self.channel_name) @@ -155,7 +123,7 @@ class NotifyConsumer(WebsocketConsumer): self.send(text_data=json.dumps(event)) -class PubSubConsumer(WebsocketConsumer): +class PubSubConsumer(BaseConsumer): def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) self.token = self.scope['url_route']['kwargs']['token'] @@ -163,9 +131,6 @@ class PubSubConsumer(WebsocketConsumer): self.p = self.rds.pubsub(ignore_subscribe_messages=True) self.p.subscribe(self.token) - def connect(self): - self.accept() - def disconnect(self, code): self.p.close() self.rds.close() diff --git a/spug_api/consumer/middleware.py b/spug_api/consumer/middleware.py deleted file mode 100644 index bba5d58..0000000 --- a/spug_api/consumer/middleware.py +++ /dev/null @@ -1,50 +0,0 @@ -# Copyright: (c) OpenSpug Organization. https://github.com/openspug/spug -# Copyright: (c) -# Released under the AGPL-3.0 License. -from django.db import close_old_connections -from channels.security.websocket import WebsocketDenier -from apps.account.models import User -from apps.setting.utils import AppSetting -from libs.utils import get_request_real_ip -from urllib.parse import parse_qs -import time - - -class AuthMiddleware: - def __init__(self, application): - self.application = application - - def __call__(self, scope): - # Make sure the scope is of type websocket - if scope["type"] != "websocket": - raise ValueError( - "You cannot use AuthMiddleware on a non-WebSocket connection" - ) - headers = dict(scope.get('headers', [])) - is_ok, message = self.verify_user(scope, headers) - if is_ok: - return self.application(scope) - else: - print(message) - return WebsocketDenier(scope) - - def get_real_ip(self, headers): - decode_headers = { - 'x-forwarded-for': headers.get(b'x-forwarded-for', b'').decode(), - 'x-real-ip': headers.get(b'x-real-ip', b'').decode() - } - return get_request_real_ip(decode_headers) - - def verify_user(self, scope, headers): - close_old_connections() - query_string = scope['query_string'].decode() - x_real_ip = self.get_real_ip(headers) - token = parse_qs(query_string).get('x-token', [''])[0] - if token and len(token) == 32: - user = User.objects.filter(access_token=token).first() - if user and user.token_expired >= time.time() and user.is_active: - if x_real_ip == user.last_ip or AppSetting.get_default('bind_ip') is False: - scope['user'] = user - return True, None - return False, f'Verify failed: {x_real_ip} <> {user.last_ip if user else None}' - return False, 'Token is invalid' diff --git a/spug_api/consumer/routing.py b/spug_api/consumer/routing.py index 7055c0a..9ef58b9 100644 --- a/spug_api/consumer/routing.py +++ b/spug_api/consumer/routing.py @@ -3,15 +3,11 @@ # Released under the AGPL-3.0 License. from django.urls import path from channels.routing import URLRouter -from consumer.middleware import AuthMiddleware from consumer.consumers import * -ws_router = AuthMiddleware( - URLRouter([ - path('ws/exec//', ExecConsumer), - path('ws/ssh//', SSHConsumer), - path('ws/subscribe//', PubSubConsumer), - path('ws///', ComConsumer), - path('ws/notify/', NotifyConsumer), - ]) -) +ws_router = URLRouter([ + path('ws/ssh//', SSHConsumer), + path('ws/subscribe//', PubSubConsumer), + path('ws///', ComConsumer), + path('ws/notify/', NotifyConsumer), +]) diff --git a/spug_api/consumer/utils.py b/spug_api/consumer/utils.py new file mode 100644 index 0000000..1313b8b --- /dev/null +++ b/spug_api/consumer/utils.py @@ -0,0 +1,42 @@ +# Copyright: (c) OpenSpug Organization. https://github.com/openspug/spug +# Copyright: (c) +# Released under the AGPL-3.0 License. +from django.db import close_old_connections +from channels.generic.websocket import WebsocketConsumer +from apps.account.models import User +from apps.setting.utils import AppSetting +from libs.utils import get_request_real_ip +from urllib.parse import parse_qs +import time + + +def get_real_ip(headers): + decode_headers = {k.decode(): v.decode() for k, v in headers} + return get_request_real_ip(decode_headers) + + +class BaseConsumer(WebsocketConsumer): + def __init__(self, *args, **kwargs): + super(BaseConsumer, self).__init__(*args, **kwargs) + self.user = None + + def close_with_message(self, content): + self.send(text_data=f'\r\n\x1b[31m{content}\x1b[0m\r\n') + self.close() + + def connect(self): + self.accept() + close_old_connections() + query_string = self.scope['query_string'].decode() + x_real_ip = get_real_ip(self.scope['headers']) + token = parse_qs(query_string).get('x-token', [''])[0] + if token and len(token) == 32: + user = User.objects.filter(access_token=token).first() + if user and user.token_expired >= time.time() and user.is_active: + if x_real_ip == user.last_ip or AppSetting.get_default('bind_ip') is False: + self.user = user + if hasattr(self, 'init'): + self.init() + return None + self.close_with_message('触发登录IP绑定安全策略,请在系统设置/安全设置中查看配置。') + self.close_with_message('用户身份验证失败,请重新登录或刷新页面。') diff --git a/spug_web/src/layout/Notification.js b/spug_web/src/layout/Notification.js index 6f25d55..24395be 100644 --- a/spug_web/src/layout/Notification.js +++ b/spug_web/src/layout/Notification.js @@ -70,11 +70,15 @@ export default function () { ws.onmessage = e => { if (e.data !== 'pong') { fetch(); - const {title, content} = JSON.parse(e.data); - const key = `open${Date.now()}`; - const description =
{content}
; - const btn = ; - notification.warning({message: title, description, btn, key, top: 64, duration: null}) + try { + const {title, content} = JSON.parse(e.data); + const key = `open${Date.now()}`; + const description =
{content}
; + const btn = ; + notification.warning({message: title, description, btn, key, top: 64, duration: null}) + } catch (e) { + + } } } } diff --git a/spug_web/src/pages/exec/task/Output.js b/spug_web/src/pages/exec/task/Output.js index 0559b59..f104346 100644 --- a/spug_web/src/pages/exec/task/Output.js +++ b/spug_web/src/pages/exec/task/Output.js @@ -16,7 +16,7 @@ import { import { FitAddon } from 'xterm-addon-fit'; import { Terminal } from 'xterm'; import style from './index.module.less'; -import { X_TOKEN } from 'libs'; +import { http, X_TOKEN } from 'libs'; import store from './store'; import gStore from 'gStore'; @@ -55,7 +55,7 @@ function OutView(props) { useEffect(() => { const protocol = window.location.protocol === 'https:' ? 'wss:' : 'ws:'; - const socket = new WebSocket(`${protocol}//${window.location.host}/api/ws/exec/${store.token}/?x-token=${X_TOKEN}`); + const socket = new WebSocket(`${protocol}//${window.location.host}/api/ws/subscribe/${store.token}/?x-token=${X_TOKEN}`); socket.onopen = () => { const message = '\r\x1b[K\x1b[36m### Waiting for scheduling ...\x1b[0m' for (let key of Object.keys(store.outputs)) { @@ -64,6 +64,7 @@ function OutView(props) { term.write(message) socket.send('ok'); fitPlugin.fit() + http.patch('/api/exec/do/', {token: store.token}) } socket.onmessage = e => { if (e.data === 'pong') { diff --git a/spug_web/src/pages/exec/task/index.js b/spug_web/src/pages/exec/task/index.js index 66f4f47..a91f9eb 100644 --- a/spug_web/src/pages/exec/task/index.js +++ b/spug_web/src/pages/exec/task/index.js @@ -15,6 +15,7 @@ import Output from './Output'; import { http, cleanCommand } from 'libs'; import moment from 'moment'; import store from './store'; +import gStore from 'gStore'; import style from './index.module.less'; function TaskIndex() { @@ -28,7 +29,7 @@ function TaskIndex() { useEffect(() => { if (!loading) { - http.get('/api/exec/history/') + http.get('/api/exec/do/') .then(res => setHistories(res)) } }, [loading]) @@ -40,6 +41,7 @@ function TaskIndex() { }, [command]) useEffect(() => { + gStore.fetchUserSettings() return () => { store.host_ids = [] if (store.showConsole) { diff --git a/spug_web/src/pages/ssh/FileManager.js b/spug_web/src/pages/ssh/FileManager.js index a7abd53..7d18355 100644 --- a/spug_web/src/pages/ssh/FileManager.js +++ b/spug_web/src/pages/ssh/FileManager.js @@ -174,7 +174,7 @@ class FileManager extends React.Component { _updatePercent = token => { const protocol = window.location.protocol === 'https:' ? 'wss:' : 'ws:'; - this.socket = new WebSocket(`${protocol}//${window.location.host}/api/ws/exec/${token}/?x-token=${X_TOKEN}`); + this.socket = new WebSocket(`${protocol}//${window.location.host}/api/ws/subscribe/${token}/?x-token=${X_TOKEN}`); this.socket.onopen = () => this.socket.send('ok'); this.socket.onmessage = e => { if (e.data === 'pong') {