From ac11b1c359fd867d52f7dd59d76ec8a58b9478cf Mon Sep 17 00:00:00 2001
From: root <email@example.com>
Date: Thu, 24 Nov 2022 19:59:21 +0000
Subject: [PATCH] Monkeypatch sshim.expect, implement basic connect_handler,
 set default to localhost only

---
 main.py        | 25 +++++++++++-----------
 sshim_patch.py | 58 ++++++++++++++++++++++++++++++++++++++++++++------
 2 files changed, 63 insertions(+), 20 deletions(-)

diff --git a/main.py b/main.py
index 8f26d36..af204f6 100644
--- a/main.py
+++ b/main.py
@@ -4,26 +4,25 @@ import sshim_patch as sshim
 import lxd_interface
 import os
 import re
+import uuid
 
 logging.basicConfig(level='DEBUG')
 logger = logging.getLogger()
 
 
-def connect_handler(script):
-    # ask the SSH client to enter a name
-    script.write('Please enter your name: ')
-
-    # match their input against a regular expression which will store the name in a capturing group called name
-    groups = script.expect(re.compile('(?P<name>.*)')).groupdict()
-
-    # log on the server-side that the user has connected
-    logger.info('%(name)s just connected', **groups)
-
-    # send a message back to the SSH client greeting it by name
-    script.writeline('Hello %(name)s!' % groups)
+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!")
 
 
-server = sshim.Server(connect_handler, address='0.0.0.0', port=3000)
+server = sshim.Server(connect_handler, address='127.0.0.1', port=3022)
 try:
     server.run()
 except KeyboardInterrupt:
diff --git a/sshim_patch.py b/sshim_patch.py
index 813c761..3b56c57 100644
--- a/sshim_patch.py
+++ b/sshim_patch.py
@@ -1,13 +1,60 @@
 from sshim import *
 import paramiko
 import os
+import six
+import codecs
+
+
+def expect(self, line, echo=True) -> str:
+    """
+        Expect a line of input from the user. If this has the `match` method, it will call it on the input and return
+        the result, otherwise it will use the equality operator, ==. Notably, if a regular expression is passed in
+        its match method will be called and the matchdata returned. This allows you to use matching groups to pull
+        out interesting data and operate on it.
+        If ``echo`` is set to False, the server will not echo the input back to the client.
+    """
+    buffer = six.BytesIO()
+
+    try:
+        while True:
+            byte = self.fileobj.read(1)
+
+            if not byte or byte == '\x04':
+                raise EOFError()
+            elif byte == b'\t':
+                pass
+            elif byte == b'\x7f':
+                if buffer.tell() > 0:
+                    self.sendall('\b \b')
+                    buffer.truncate(buffer.tell() - 1)
+            elif byte == b'\x1b' and self.fileobj.read(1) == b'[':
+                command = self.fileobj.read(1)
+                if hasattr(self.delegate, 'cursor'):
+                    self.delegate.cursor(command)
+            elif byte in (b'\r', b'\n'):
+                break
+            else:
+                buffer.write(byte)
+                if echo:
+                    self.sendall(byte)
+
+        if echo:
+            self.sendall('\r\n')
+
+        return codecs.decode(buffer.getvalue(), self.encoding)
+
+    except Exception:
+        raise
 
 
 def check_auth_none(self, username):
-    return paramiko.AUTH_PARTIALLY_SUCCESSFUL
+    if username == os.environ["ssh-username"]:
+        return paramiko.AUTH_PARTIALLY_SUCCESSFUL
+    return paramiko.AUTH_FAILED
 
 
 def check_auth_password(self, username, password):
+    print(os.environ["ssh-username"], os.environ["ssh-password"])
     if username == os.environ["ssh-username"] and password == os.environ["ssh-password"]:
         return paramiko.AUTH_SUCCESSFUL
     return paramiko.AUTH_FAILED
@@ -17,13 +64,10 @@ def check_auth_publickey(self, username, key):
     return paramiko.AUTH_FAILED
 
 
-def enable_auth_gssapi(self):
-    return paramiko.AUTH_FAILED
-
+Script.expect = expect
 
 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 = enable_auth_gssapi
-Handler.get_allowed_auths = paramiko.server.ServerInterface.get_allowed_auths
-
+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