From c87598f73781406458339438106713f07d7d12b0 Mon Sep 17 00:00:00 2001
From: root <email@example.com>
Date: Fri, 25 Nov 2022 15:41:02 +0000
Subject: [PATCH] Move connect_handler logic to sshim_patch

---
 main.py        | 12 +++---------
 sshim_patch.py | 34 +++++++++++++++++++++++++++++++++-
 2 files changed, 36 insertions(+), 10 deletions(-)

diff --git a/main.py b/main.py
index af204f6..e192c7f 100644
--- a/main.py
+++ b/main.py
@@ -1,4 +1,6 @@
 import logging
+import time
+
 import paramiko
 import sshim_patch as sshim
 import lxd_interface
@@ -11,15 +13,7 @@ logger = logging.getLogger()
 
 
 def connect_handler(script: sshim.Script):
-    instance_name = "instance-" + str(uuid.uuid4())
-    lxd_interface.create_instance(instance_name)
-    with paramiko.ProxyCommand(command_line=f'lxc exec {instance_name} -- /bin/bash') as proxy_command:
-        script.writeline(instance_name)
-        while True:
-            input_command = script.expect(None, echo=True) # TODO: change to false
-            proxy_command.send(input_command.encode())
-            script.sendall(proxy_command.recv(100))  # TODO: fix
-            script.writeline("Sent!")
+    pass
 
 
 server = sshim.Server(connect_handler, address='127.0.0.1', port=3022)
diff --git a/sshim_patch.py b/sshim_patch.py
index 3b56c57..8ea2710 100644
--- a/sshim_patch.py
+++ b/sshim_patch.py
@@ -3,6 +3,13 @@ import paramiko
 import os
 import six
 import codecs
+import uuid
+import lxd_interface
+import threading
+import logging
+
+logging.basicConfig(level='DEBUG')
+logger = logging.getLogger()
 
 
 def expect(self, line, echo=True) -> str:
@@ -47,6 +54,13 @@ def expect(self, line, echo=True) -> str:
         raise
 
 
+def check_channel_shell_request(self, channel):
+    logger.debug(channel)
+    Runner(self, channel).start()
+
+    return True
+
+
 def check_auth_none(self, username):
     if username == os.environ["ssh-username"]:
         return paramiko.AUTH_PARTIALLY_SUCCESSFUL
@@ -64,10 +78,28 @@ def check_auth_publickey(self, username, key):
     return paramiko.AUTH_FAILED
 
 
+class Runner(threading.Thread):
+    def __init__(self, client, channel: paramiko.Channel):
+        threading.Thread.__init__(self, name='sshim.Runner(%s)' % channel.chanid)
+        self.instance_name = "instance-" + str(uuid.uuid4())
+        self.daemon = True
+        self.client = client
+        self.channel = channel
+        self.channel.settimeout(None)
+
+    def run(self) -> None:
+        lxd_interface.create_instance(self.instance_name)
+
+        with paramiko.ProxyCommand(command_line=f'lxc exec {self.instance_name} -- /bin/bash') as proxy_command:
+            self.channel.recv = proxy_command.recv
+            self.channel.send = proxy_command.send
+
+
 Script.expect = expect
 
+Handler.check_channel_shell_request = check_channel_shell_request
 Handler.check_auth_none = check_auth_none
 Handler.check_auth_password = check_auth_password
 Handler.check_auth_publickey = check_auth_publickey
 Handler.enable_auth_gssapi = paramiko.server.ServerInterface.enable_auth_gssapi
-Handler.get_allowed_auths = paramiko.server.ServerInterface.get_allowed_auths
\ No newline at end of file
+Handler.get_allowed_auths = paramiko.server.ServerInterface.get_allowed_auths