import paramiko
import select
import socket
import threading
import subprocess
import os
import pty


class CustomSSHServer(paramiko.ServerInterface):
    def check_auth_password(self, username, password):
        # Implement your authentication logic here
        return paramiko.AUTH_SUCCESSFUL

    def check_channel_request(self, kind, chanid):
        return paramiko.OPEN_SUCCEEDED


def get_vm_port(username):
    return 9002


def handle_client(client_socket, addr):
    transport = paramiko.Transport(client_socket)
    # transport.add_server_key(paramiko.RSAKey(filename='server.key'))
    transport.add_server_key(paramiko.Ed25519Key.from_private_key_file('daemon.key'))
    server = CustomSSHServer()
    transport.start_server(server=server)

    channel = transport.accept(20)
    if channel is None:
        print('*** No channel.')
        return

    username = transport.get_username()
    vm_port = get_vm_port(username)
    master, slave = pty.openpty()
    channel.get_pty()
    channel.invoke_shell()

    ssh = subprocess.Popen(
        ["ssh", "-p", "9002", "root@localhost"],
        shell=True,
        stdin=slave,
        stdout=slave,
        stderr=slave,
        close_fds=True
    )
    #ssh.wait()

    # Forward data between the channel and the process
    while True:
        r, w, e = select.select([channel, master], [], [])
        if channel in r:
            data = channel.recv(1024)
            if not data:
                break
            os.write(master, data)
        if master in r:
            data = os.read(master, 1024)
            if not data:
                break
            channel.send(data)

    channel.close()
    ssh.terminate()


def main():
    server_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    server_socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    server_socket.bind(('', 2222))
    server_socket.listen(100)

    while True:
        client_socket, addr = server_socket.accept()
        threading.Thread(target=handle_client, args=(client_socket, addr)).start()


if __name__ == '__main__':
    main()