A api update

spug_api/apps/consumer/consumers.py

@@ -1,22 +0,0 @@
-from channels.generic.websocket import WebsocketConsumer
-from django_redis import get_redis_connection
-
-
-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 receive(self, **kwargs):
-        response = self.rds.blpop(self.token, timeout=5)
-        while response:
-            self.send(text_data=response[1].decode())
-            response = self.rds.blpop(self.token, timeout=5)
-        self.send(text_data='pong')

spug_api/apps/host/urls.py

@@ -1,7 +1,8 @@
-from django.conf.urls import url
+from django.urls import path
 
 from .views import *
 
 urlpatterns = [
-    url(r'^$', HostView.as_view()),
+    path('', HostView.as_view()),
+    path('ssh/<int:h_id>/', web_ssh),
 ]

spug_api/apps/host/views.py

@@ -1,4 +1,6 @@
 from django.views.generic import View
+from django.shortcuts import render
+from django.http.response import HttpResponseBadRequest
 from libs import json_response, JsonParser, Argument
 from apps.setting.utils import AppSetting
 from apps.host.models import Host
@@ -46,6 +48,14 @@ class HostView(View):
         return json_response(error=error)
 
 
+def web_ssh(request, h_id):
+    host = Host.objects.filter(pk=h_id).first()
+    if not host:
+        return HttpResponseBadRequest('unknown host')
+    context = {'id': h_id, 'title': host.name, 'token': request.user.access_token}
+    return render(request, 'web_ssh.html', context)
+
+
 def valid_ssh(hostname, port, username, password):
     try:
         private_key = AppSetting.get('private_key')

spug_api/consumer/consumers.py

@@ -0,0 +1,77 @@
+from channels.generic.websocket import WebsocketConsumer
+from django_redis import get_redis_connection
+from apps.setting.utils import AppSetting
+from apps.host.models import Host
+from threading import Thread
+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 receive(self, **kwargs):
+        response = self.rds.blpop(self.token, timeout=5)
+        while response:
+            self.send(text_data=response[1].decode())
+            response = self.rds.blpop(self.token, timeout=5)
+        self.send(text_data='pong')
+
+
+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.chan = None
+        self.ssh = None
+
+    def loop_read(self):
+        while True:
+            data = self.chan.recv(32 * 1024)
+            # print('read: {!r}'.format(data))
+            if not data:
+                self.close(3333)
+                break
+            self.send(bytes_data=data)
+
+    def receive(self, text_data=None, bytes_data=None):
+        data = text_data or bytes_data
+        if data:
+            data = json.loads(data)
+            # print('write: {!r}'.format(data))
+            resize = data.get('resize')
+            if resize and len(resize) == 2:
+                self.chan.resize_pty(*resize)
+            else:
+                self.chan.send(data['data'])
+
+    def disconnect(self, code):
+        self.chan.close()
+        self.ssh.close()
+        # print('Connection close')
+
+    def connect(self):
+        self.accept()
+        self.send(bytes_data=b'Connecting ...\r\n')
+        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(AppSetting.get('private_key')).get_client()
+        except Exception as e:
+            self.send(bytes_data=f'Exception: {e}\r\n'.encode())
+            self.close()
+        self.chan = self.ssh.invoke_shell(term='xterm')
+        self.chan.transport.set_keepalive(30)
+        Thread(target=self.loop_read).start()

spug_api/consumer/routing.py

@@ -3,4 +3,5 @@ from .consumers import *
 
 websocket_urlpatterns = [
     path('ws/exec/<str:token>/', ExecConsumer),
+    path('ws/ssh/<str:token>/<int:id>/', SSHConsumer),
 ]

spug_api/libs/ssh.py

@@ -38,6 +38,14 @@ class SSH:
         with self:
             return True
 
+    def get_client(self):
+        if self.client is not None:
+            return self.client
+        self.client = SSHClient()
+        self.client.set_missing_host_key_policy(AutoAddPolicy)
+        self.client.connect(**self.arguments)
+        return self.client
+
     def exec_command(self, command, timeout=1800, environment=None):
         with self as cli:
             chan = cli.get_transport().open_session()
@@ -67,11 +75,7 @@ class SSH:
     def __enter__(self):
         if self.client is not None:
             raise RuntimeError('Already connected')
-        client = SSHClient()
-        client.set_missing_host_key_policy(AutoAddPolicy)
-        client.connect(**self.arguments)
-        self.client = client
-        return self.client
+        return self.get_client()
 
     def __exit__(self, exc_type, exc_val, exc_tb):
         self.client.close()

spug_api/spug/routing.py

@@ -1,5 +1,5 @@
 from channels.routing import ProtocolTypeRouter, ChannelNameRouter, URLRouter
-from apps.consumer import routing, executors
+from consumer import routing, executors
 
 application = ProtocolTypeRouter({
     'channel': ChannelNameRouter({

spug_api/spug/settings.py

@@ -82,6 +82,14 @@ CHANNEL_LAYERS = {
     },
 }
 
+TEMPLATES = [
+    {
+        'BACKEND': 'django.template.backends.django.DjangoTemplates',
+        'DIRS': [os.path.join(BASE_DIR, 'templates')],
+        'APP_DIRS': False,
+    },
+]
+
 SCHEDULE_KEY = 'spug:schedule'
 MONITOR_KEY = 'spug:monitor'
 

spug_api/templates/web_ssh.html

@@ -0,0 +1,22 @@
+<!DOCTYPE html>
+<html lang="en">
+<head>
+  <meta charset="UTF-8">
+  <title>{{ title }}</title>
+  <link href="http://cdn.qbangmang.com/spug/xterm.min.css" rel="stylesheet" type="text/css"/>
+  <style>
+    body {
+      margin: 0;
+    }
+  </style>
+</head>
+<body>
+<div id="terminal"></div>
+<script src="http://cdn.qbangmang.com/spug/xterm.min.js"></script>
+<script src="http://cdn.qbangmang.com/spug/xterm-addon-fit.min.js"></script>
+<script src="http://localhost:63342/spug_v2/main.js"></script>
+<script>
+  run('{{ id }}', '{{ token }}')
+</script>
+</body>
+</html>