U improve websocket security

spug_api/consumer/consumers.py

@@ -3,11 +3,9 @@
 # Released under the AGPL-3.0 License.
 from channels.generic.websocket import WebsocketConsumer
 from django_redis import get_redis_connection
-from apps.account.models import User
 from apps.host.models import Host
 from threading import Thread
 import json
-import time
 
 
 class ExecConsumer(WebsocketConsumer):
@@ -38,9 +36,8 @@ class ExecConsumer(WebsocketConsumer):
 class SSHConsumer(WebsocketConsumer):
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
-        kwargs = self.scope['url_route']['kwargs']
-        self.token = kwargs['token']
-        self.id = kwargs['id']
+        self.user = self.scope['user']
+        self.id = self.scope['url_route']['kwargs']['id']
         self.chan = None
         self.ssh = None
 
@@ -70,8 +67,7 @@ class SSHConsumer(WebsocketConsumer):
         # print('Connection close')
 
     def connect(self):
-        user = User.objects.filter(access_token=self.token).first()
-        if user and user.token_expired >= time.time() and user.is_active and user.has_host_perm(self.id):
+        if self.user.has_host_perm(self.id):
             self.accept()
             self._init()
         else:

spug_api/consumer/routing.py

@@ -2,9 +2,13 @@
 # Copyright: (c) <spug.dev@gmail.com>
 # Released under the AGPL-3.0 License.
 from django.urls import path
-from .consumers import *
+from channels.routing import URLRouter
+from consumer.middleware import AuthMiddleware
+from consumer.consumers import *
 
-websocket_urlpatterns = [
+ws_router = AuthMiddleware(
+    URLRouter([
         path('ws/exec/<str:token>/', ExecConsumer),
-    path('ws/ssh/<str:token>/<int:id>/', SSHConsumer),
-]
+        path('ws/ssh/<int:id>/', SSHConsumer),
+    ])
+)

spug_api/libs/utils.py

@@ -105,10 +105,7 @@ def generate_random_str(length: int = 4, is_digits: bool = True) -> str:
 
 
 def get_request_real_ip(headers: dict):
-    x_real_ip = headers.get('x-real-ip')
+    x_real_ip = headers.get('x-forwarded-for')
     if not x_real_ip:
-        x_forwarded_for = headers.get('x-forwarded-for')
-        if not x_forwarded_for:
-            return ''
-        x_real_ip = x_forwarded_for.split(',')[0]
-    return x_real_ip
+        x_real_ip = headers.get('x-real-ip', '')
+    return x_real_ip.split(',')[0]

spug_api/spug/routing.py

@@ -1,14 +1,12 @@
 # Copyright: (c) OpenSpug Organization. https://github.com/openspug/spug
 # Copyright: (c) <spug.dev@gmail.com>
 # Released under the AGPL-3.0 License.
-from channels.routing import ProtocolTypeRouter, ChannelNameRouter, URLRouter
+from channels.routing import ProtocolTypeRouter, ChannelNameRouter
 from consumer import routing, executors
 
 application = ProtocolTypeRouter({
     'channel': ChannelNameRouter({
         'ssh_exec': executors.SSHExecutor,
     }),
-    'websocket': URLRouter(
-        routing.websocket_urlpatterns
-    )
+    'websocket': routing.ws_router
 })

spug_web/src/libs/functools.js

@@ -9,7 +9,10 @@ let Permission = {
   permissions: []
 };
 
+export let X_TOKEN;
+
 export function updatePermissions() {
+  X_TOKEN = localStorage.getItem('token');
   Permission.isSuper = localStorage.getItem('is_supper') === 'true';
   Permission.hostPerms = JSON.parse(localStorage.getItem('host_perms') || '[]');
   Permission.permissions = JSON.parse(localStorage.getItem('permissions') || '[]');

spug_web/src/libs/http.js

@@ -5,7 +5,8 @@
  */
 import http from 'axios'
 import history from './history'
-import {message} from 'antd';
+import { X_TOKEN } from './functools';
+import { message } from 'antd';
 
 // response处理
 function handleResponse(response) {
@@ -40,7 +41,7 @@ function handleResponse(response) {
 http.interceptors.request.use(request => {
   request.isInternal = request.url.startsWith('/api/');
   if (request.isInternal) {
-    request.headers['X-Token'] = localStorage.getItem('token')
+    request.headers['X-Token'] = X_TOKEN
   }
   request.timeout = request.timeout || 30000;
   return request;

spug_web/src/pages/deploy/do/Ext1Index.js

@@ -6,10 +6,9 @@
 import React from 'react';
 import { observer } from 'mobx-react';
 import { Steps, Collapse, PageHeader, Spin, Tag, Button, Icon } from 'antd';
-import http from 'libs/http';
+import { http, history, X_TOKEN } from 'libs';
 import { AuthDiv } from 'components';
 import OutView from './OutView';
-import history from 'libs/history';
 import styles from './index.module.css';
 import store from './store';
 import lds from 'lodash';
@@ -79,7 +78,7 @@ class Ext1Index extends React.Component {
         store.request.status = '2';
         store.outputs = outputs;
         const protocol = window.location.protocol === 'https:' ? 'wss:' : 'ws:';
-        this.socket = new WebSocket(`${protocol}//${window.location.host}/api/ws/exec/${token}/`);
+        this.socket = new WebSocket(`${protocol}//${window.location.host}/api/ws/exec/${token}/?x-token=${X_TOKEN}`);
         this.socket.onopen = () => {
           this.socket.send('ok');
         };

spug_web/src/pages/deploy/do/Ext2Index.js

@@ -6,10 +6,9 @@
 import React from 'react';
 import { observer } from 'mobx-react';
 import { Steps, Collapse, PageHeader, Spin, Tag, Button, Icon } from 'antd';
-import http from 'libs/http';
+import { http, history, X_TOKEN } from 'libs';
 import { AuthDiv } from 'components';
 import OutView from './OutView';
-import history from 'libs/history';
 import styles from './index.module.css';
 import store from './store';
 import lds from 'lodash';
@@ -80,7 +79,7 @@ class Ext1Index extends React.Component {
         store.request.status = '2';
         store.outputs = outputs;
         const protocol = window.location.protocol === 'https:' ? 'wss:' : 'ws:';
-        this.socket = new WebSocket(`${protocol}//${window.location.host}/api/ws/exec/${token}/`);
+        this.socket = new WebSocket(`${protocol}//${window.location.host}/api/ws/exec/${token}/?x-token=${X_TOKEN}`);
         this.socket.onopen = () => {
           this.socket.send('ok');
         };

spug_web/src/pages/deploy/request/Ext2Form.js

@@ -7,7 +7,7 @@ import React from 'react';
 import { observer } from 'mobx-react';
 import { Modal, Form, Input, Tag, Upload, message, Button } from 'antd';
 import hostStore from 'pages/host/store';
-import http from 'libs/http';
+import { http, X_TOKEN } from 'libs';
 import store from './store';
 import lds from 'lodash';
 
@@ -15,7 +15,6 @@ import lds from 'lodash';
 class Ext2Form extends React.Component {
   constructor(props) {
     super(props);
-    this.token = localStorage.getItem('token');
     this.state = {
       loading: false,
       uploading: false,
@@ -112,7 +111,7 @@ class Ext2Form extends React.Component {
             )}
           </Form.Item>
           <Form.Item label="上传数据" help="通过数据传输动作来使用上传的文件。">
-            <Upload name="file" fileList={fileList} headers={{'X-Token': this.token}} beforeUpload={this.handleUpload}
+            <Upload name="file" fileList={fileList} headers={{'X-Token': X_TOKEN}} beforeUpload={this.handleUpload}
                     data={{deploy_id: info.deploy_id}} onChange={this.handleUploadChange}>
               {fileList.length === 0 ? <Button loading={uploading} icon="upload">点击上传</Button> : null}
             </Upload>

spug_web/src/pages/exec/task/ExecConsole.js

@@ -6,6 +6,7 @@
 import React from 'react';
 import { observer } from 'mobx-react';
 import { Modal, Collapse, Tooltip, Icon } from 'antd';
+import { X_TOKEN } from 'libs';
 import OutView from './OutView';
 import styles from './index.module.css';
 import store from './store';
@@ -24,7 +25,7 @@ class ExecConsole extends React.Component {
 
   componentDidMount() {
     const protocol = window.location.protocol === 'https:' ? 'wss:' : 'ws:';
-    this.socket = new WebSocket(`${protocol}//${window.location.host}/api/ws/exec/${store.token}/`);
+    this.socket = new WebSocket(`${protocol}//${window.location.host}/api/ws/exec/${store.token}/?x-token=${X_TOKEN}`);
     this.socket.onopen = () => {
       this.socket.send('ok');
       for (let item of Object.values(store.outputs)) {

spug_web/src/pages/host/Form.js

@@ -6,14 +6,13 @@
 import React from 'react';
 import { observer } from 'mobx-react';
 import { Modal, Form, Input, Select, Col, Button, Upload, message } from 'antd';
-import http from 'libs/http';
+import { http, X_TOKEN } from 'libs';
 import store from './store';
 
 @observer
 class ComForm extends React.Component {
   constructor(props) {
     super(props);
-    this.token = localStorage.getItem('token');
     this.state = {
       loading: false,
       uploading: false,
@@ -200,7 +199,7 @@ class ComForm extends React.Component {
             </Form.Item>
           </Form.Item>
           <Form.Item label="独立密钥" extra="默认使用全局密钥，如果上传了独立密钥则优先使用该密钥。">
-            <Upload name="file" fileList={fileList} headers={{'X-Token': this.token}} beforeUpload={this.handleUpload}
+            <Upload name="file" fileList={fileList} headers={{'X-Token': X_TOKEN}} beforeUpload={this.handleUpload}
                     onChange={this.handleUploadChange}>
               {fileList.length === 0 ? <Button loading={uploading} icon="upload">点击上传</Button> : null}
             </Upload>

spug_web/src/pages/ssh/FileManager.js

@@ -5,7 +5,7 @@
  */
 import React from 'react';
 import { Drawer, Breadcrumb, Table, Icon, Divider, Switch, Button, Progress, Modal, message } from 'antd';
-import { http, uniqueId } from 'libs';
+import { http, uniqueId, X_TOKEN } from 'libs';
 import lds from 'lodash';
 import styles from './index.module.css'
 
@@ -129,7 +129,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}/`);
+    this.socket = new WebSocket(`${protocol}//${window.location.host}/api/ws/exec/${token}/?x-token=${X_TOKEN}`);
     this.socket.onopen = () => this.socket.send('ok');
     this.socket.onmessage = e => {
       if (e.data === 'pong') {
@@ -145,9 +145,8 @@ class FileManager extends React.Component {
 
   handleDownload = (name) => {
     const file = `/${this.state.pwd.join('/')}/${name}`;
-    const token = localStorage.getItem('token');
     const link = document.createElement('a');
-    link.href = `/api/file/object/?id=${this.id}&file=${file}&x-token=${token}`;
+    link.href = `/api/file/object/?id=${this.id}&file=${file}&x-token=${X_TOKEN}`;
     document.body.appendChild(link);
     const evt = document.createEvent("MouseEvents");
     evt.initEvent("click", false, false);

spug_web/src/pages/ssh/index.js

@@ -9,7 +9,7 @@ import { AuthDiv } from 'components';
 import { Terminal } from 'xterm';
 import { FitAddon } from 'xterm-addon-fit';
 import FileManager from './FileManager';
-import { http } from 'libs';
+import { http, X_TOKEN } from 'libs';
 import 'xterm/css/xterm.css';
 import styles from './index.module.css';
 
@@ -18,7 +18,6 @@ class WebSSH extends React.Component {
   constructor(props) {
     super(props);
     this.id = props.match.params.id;
-    this.token = localStorage.getItem('token');
     this.socket = null;
     this.term = new Terminal();
     this.container = null;
@@ -37,7 +36,7 @@ class WebSSH extends React.Component {
     const fitPlugin = new FitAddon();
     this.term.loadAddon(fitPlugin);
     const protocol = window.location.protocol === 'https:' ? 'wss:' : 'ws:';
-    this.socket = new WebSocket(`${protocol}//${window.location.host}/api/ws/ssh/${this.token}/${this.id}/`);
+    this.socket = new WebSocket(`${protocol}//${window.location.host}/api/ws/ssh/${this.id}/?x-token=${X_TOKEN}`);
     this.socket.onmessage = e => this._read_as_text(e.data);
     this.socket.onopen = () => {
       this.term.open(this.container);

spug_web/src/setupProxy.js

@@ -10,7 +10,7 @@ module.exports = function (app) {
     target: 'http://127.0.0.1:8000',
     changeOrigin: true,
     ws: true,
-    headers: {'X-Real-IP': '127.0.0.1'},
+    headers: {'X-Real-IP': '1.1.1.1'},
     pathRewrite: {
       '^/api': ''
     }