Add useful script

lxd_interface.py

@@ -12,11 +12,10 @@ def create_instance(container_name: str, instance_password: str):
 
     instance = lxd_client.instances.create(config, wait=True)
     instance.start(wait=True)
+    setup_ssh(container_name, instance_password)
     while type(ipaddress.ip_address(instance.state().network['eth0']['addresses'][0]['address'])) != ipaddress.IPv4Address:
         time.sleep(0.1)
 
-    setup_ssh(container_name, instance_password)
-
     return instance.state().network['eth0']['addresses'][0]
 
 
@@ -53,3 +52,27 @@ def setup_ssh(container_name: str, instance_password: str):
     execute_command(container_name, ["passwd", "root"], stdin_payload=f"{instance_password}\n{instance_password}")
 
     return True
+
+
+def get_networking(container_name: str):
+    instance = lxd_client.instances.get(container_name)
+
+    while instance.state().network is None or \
+            type(ipaddress.ip_address(instance.state().network['eth0']['addresses'][0]['address'])) != ipaddress.IPv4Address:
+        time.sleep(0.1)
+
+    return instance.state().network['eth0']['addresses'][0]
+
+
+def set_description(container_name: str, new_value: str):
+    instance = lxd_client.instances.get(container_name)
+    instance.description = new_value
+    instance.save()
+
+    return True
+
+
+def get_description(container_name: str) -> str:
+    instance = lxd_client.instances.get(container_name)
+
+    return instance.description

main.py

@@ -15,7 +15,7 @@ def connect_handler(script: sshim.Script):
     pass
 
 
-server = sshim.Server(connect_handler, address='127.0.0.1', port=3022)
+server = sshim.Server(connect_handler, address='0.0.0.0', port=3022)
 try:
     server.run()
 except KeyboardInterrupt:

sshim_patch.py

@@ -1,4 +1,5 @@
 from sshim import *
+import pylxd
 import paramiko
 import os
 import uuid
@@ -7,6 +8,8 @@ import threading
 import logging
 import select
 import time
+import ipaddress
+import secrets
 import inspect
 
 logger = logging.getLogger(__name__)
@@ -21,30 +24,38 @@ def check_channel_request(self, kind, channel_id):
 
 def check_channel_shell_request(self, channel):
     logger.debug("Check channel shell request: %s" % channel.get_id())
-    self.runner.set_shell_channel(channel)
+    Runner(self, self.username, 'shell', channel).start()
+
+    return True
+
+
+def check_channel_exec_request(self, channel):
+    logger.debug("Check channel exec request: %s" % channel.get_id())
+    Runner(self, self.username, 'exec', channel).start()
 
     return True
 
 
 def check_channel_subsystem_request(self, channel, name):
     if name == 'sftp':
-        self.runner.set_sftp_channel(channel)
+        Runner(self, self.username, 'sftp', channel).start()
         return True
     else:
         return False
 
 
 def check_auth_none(self, username):
-    if username == os.environ["SSH_USERNAME"]:
-        return paramiko.AUTH_PARTIALLY_SUCCESSFUL
     return paramiko.AUTH_FAILED
 
 
 def check_auth_password(self, username, password):
     logger.debug(f"{username} just tried to connect")
-    if username == os.environ["SSH_USERNAME"] and password == os.environ["SSH_PASSWORD"]:
-        self.runner = Runner(self, self.transport)
-        self.runner.start()
+    # ensure that the connection is made from a local ip
+    if ipaddress.ip_address(self.address).is_private is not True:
+        return paramiko.AUTH_FAILED
+    if secrets.compare_digest(password, os.environ["SSH_PASSWORD"]):
+        self.username = username
+        Runner(self, self.username).start()
         return paramiko.AUTH_SUCCESSFUL
     return paramiko.AUTH_FAILED
 
@@ -54,73 +65,85 @@ def check_auth_publickey(self, username, key):
 
 
 class Runner(threading.Thread):
-    def __init__(self, client, transport: paramiko.Transport):
-        self.instance_name = "instance-" + str(uuid.uuid4())
-        threading.Thread.__init__(self, name=f'sshim.Runner {self.instance_name}')
-        self.instance_password = str(uuid.uuid4())  # TODO: secure password generation
+    def __init__(self, client, username: str, channel_type: str = None, channel: paramiko.Channel = None):
+        self.instance_name = "instance-" + username
+        self.runner_identifier = "runner-" + str(uuid.uuid4())
+        threading.Thread.__init__(self, name=f'sshim.Runner {self.instance_name} {self.runner_identifier}')
+        self.instance_password = self.instance_name  # TODO: fix - VERY INSECURE!
         self.daemon = True
         self.client = client
-        self.transport = transport
-        # self.transport.set_subsystem_handler('sftp', handler=paramiko.SFTPServer)
-        self.shell_channel = None
-        self.sftp_channel = None
+        self.channel_type = channel_type
+        self.channel = channel
 
     def run(self) -> None:
-        vm_ip = lxd_interface.create_instance(self.instance_name, self.instance_password)['address']
+        try:
+            vm_ip = lxd_interface.create_instance(self.instance_name, self.instance_password)['address']
+        except pylxd.exceptions.LXDAPIException as e:
+            logger.debug(e)
+            vm_ip = lxd_interface.get_networking(self.instance_name)['address']
+
+        if self.channel is not None:
+            self.channel.get_transport().set_keepalive(5)  # TODO: make config option
 
         with paramiko.SSHClient() as ssh_client:
+            # wait for instance to start if it hasn't started yet
+            is_not_int = True
+            while is_not_int and self.channel is not None:
+                try:
+                    int(lxd_interface.get_description(self.instance_name))
+                    is_not_int = False
+                except ValueError as e:
+                    logger.debug(e)
+                    time.sleep(0.2)
+
             ssh_client.set_missing_host_key_policy(paramiko.WarningPolicy)
             ssh_client.connect(vm_ip, username='root', password=self.instance_password)
-            client_shell_channel = ssh_client.invoke_shell()
-            client_sftp_channel = ssh_client.open_sftp().get_channel()
+            if self.channel_type == "shell" or self.channel_type == "exec":
+                client_channel = ssh_client.invoke_shell()
+            elif self.channel_type == "sftp":
+                client_channel = ssh_client.open_sftp().get_channel()
 
+            last_save_time = round(time.time())
+            lxd_interface.set_description(self.instance_name, str(last_save_time))
+            forward_channel_return = True
+            while forward_channel_return is True:
+                current_time_rounded = round(time.time())
+                if current_time_rounded != last_save_time:
+                    lxd_interface.set_description(self.instance_name, str(current_time_rounded))
+                    last_save_time = current_time_rounded
+
+                if "client_channel" in locals():
+                    forward_channel_return = self.forward_channel(client_channel)
+
+            exit_time = round(time.time())
+            time.sleep(10)
             while True:
-                if self.shell_channel is not None:
-                    r, w, e = select.select([client_shell_channel, self.shell_channel], [], [])
-                    if self.shell_channel in r:
-                        x = self.shell_channel.recv(1024)
-                        if len(x) == 0:
-                            self.shell_channel.close()
-                            self.shell_channel = None
-                            continue
-                        client_shell_channel.send(x)
-                    if client_shell_channel in r:
-                        x = client_shell_channel.recv(1024)
-                        if len(x) == 0:
-                            self.shell_channel.close()
-                            self.shell_channel = None
-                            continue
-                        self.shell_channel.send(x)
+                last_run_time = int(lxd_interface.get_description(self.instance_name))
 
-                if self.sftp_channel is not None:  # TODO: move this to function
-                    r, w, e = select.select([client_sftp_channel, self.sftp_channel], [], [])
-                    if self.sftp_channel in r:
-                        x = self.sftp_channel.recv(1024)
-                        if len(x) == 0:
-                            self.sftp_channel.close()
-                            self.sftp_channel = None
-                            continue
-                        client_sftp_channel.send(x)
-                    if client_sftp_channel in r:
-                        x = client_sftp_channel.recv(1024)
-                        if len(x) == 0:
-                            self.sftp_channel.close()
-                            self.sftp_channel = None
-                            continue
-                        self.sftp_channel.send(x)
-
-                if self.transport.is_active() is False:
+                if exit_time < last_run_time:
+                    break
+                elif round(time.time()) > (last_run_time + 15):
+                    lxd_interface.destroy_instance(self.instance_name)
                     break
 
-            lxd_interface.destroy_instance(self.instance_name)
-
-    def set_shell_channel(self, channel):
-        self.shell_channel = channel
-        self.shell_channel.settimeout(None)
-
-    def set_sftp_channel(self, channel):
-        self.sftp_channel = channel
-        self.sftp_channel.settimeout(None)
+    def forward_channel(self, client_channel) -> bool:
+        if self.channel is None:
+            return False
+        else:
+            r, w, e = select.select([client_channel, self.channel], [], [])
+            if self.channel in r:
+                x = self.channel.recv(1024)
+                if len(x) == 0:
+                    self.channel.close()
+                    return False
+                client_channel.send(x)
+            if client_channel in r:
+                x = client_channel.recv(1024)
+                if len(x) == 0:
+                    self.channel.close()
+                    return False
+                self.channel.send(x)
+            return True
 
 
 Handler.check_channel_request = check_channel_request

useful/delete-all.sh

@@ -0,0 +1,2 @@
+#!/bin/bash
+lxc delete $(lxc list -c n --format csv) --force --verbose