task2/main.py - master - georgiy/work

task2/main.py
250 строк9 KB
Georgiy Yankovskiy
task2
17 ноя 2024, 23:50
17 ноя 2024, 23:5095703e6
import io
import docker
import os
import paramiko
import socket
import threading
import traceback
import resource
import subprocess
import logging
import select
from cmd import Cmd
from abc import ABC, abstractmethod
from sys import platform

from paramiko.channel import ChannelFile, Channel

logging.basicConfig(filename='terminal.log', level=logging.INFO, format='%(asctime)s %(levelname)s: %(message)s')


def log_command(command, username):
    logging.info(f"{username}: {command.rstrip()}")


class ServerBase(ABC):

    def __init__(self):
        self._is_running = threading.Event()
        self._socket = None
        self.client_shell = None
        self._listen_thread = None

    def start(self, address='127.0.0.1', port=2222, timeout=1):
        if not self._is_running.is_set():
            self._is_running.set()

            self._socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            self._socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, True)

            if platform == "linux" or platform == "linux2":
                self._socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEPORT, True)

            # self._socket.listen(2)
            # self._socket.setblocking(False)
            self._socket.settimeout(timeout)
            self._socket.bind((address, port))

            self._listen_thread = threading.Thread(target=self._listen)
            self._listen_thread.start()

    def stop(self):
        if self._is_running.is_set():
            self._is_running.clear()
            self._listen_thread.join()
            self._socket.close()

    def _listen(self):
        while self._is_running.is_set():
            try:
                self._socket.listen()
                client, addr = self._socket.accept()
                self.connection_function(client)
            except socket.timeout:
                pass

    @abstractmethod
    def connection_function(self, client):
        pass


class SshServerInterface(paramiko.ServerInterface):
    def check_channel_request(self, kind, chanid):
        if kind == "session":
            return paramiko.OPEN_SUCCEEDED
        return paramiko.OPEN_FAILED_ADMINISTRATIVELY_PROHIBITED

    def check_channel_pty_request(self, channel, term, width, height, pixelwidth, pixelheight, modes):
        return True

    def check_channel_shell_request(self, channel):
        return True

    def check_auth_password(self, username, password):
        if (((username == 'admin') or (username == 'admin1') or (username == 'admin2') or (username == 'admin3'))
                and (password == 'password')):
            return paramiko.AUTH_SUCCESSFUL
        return paramiko.AUTH_FAILED

    def get_banner(self):
        return 'Custom SSH Server\r\n', 'ru-RU'


class SshServer(ServerBase):
    exit = False

    def __init__(self, host_key_file, host_key_file_password=None):
        super(SshServer, self).__init__()

        self._host_key = paramiko.RSAKey.from_private_key_file(host_key_file, host_key_file_password)
        self.docker_client = docker.from_env()

    def connection_function(self, client):
        try:
            session = paramiko.Transport(client)
            session.add_server_key(self._host_key)

            srv = SshServerInterface()

            try:
                session.start_server(server=srv)
            except paramiko.SSHException:
                return

            channel = session.accept()
            stdio = channel.makefile('rwU')

            self.command_shell(channel, stdio, 'admin')

            session.close()
        except Exception:
            pass

    def command_shell(self, channel: Channel, stdio: ChannelFile, username):
        client = docker.from_env()
        tag_name = f'ssh-custom-shell-for-{username}'

        try:
            # Build the Docker image (only needed the first time)
            image, logs = client.images.build(
                dockerfile='Dockerfile',
                path='./shell',
                tag=tag_name
            )
            print("Build logs:\n", logs)  # Uncomment to see build logs

            container = client.containers.run(tag_name,
                                              detach=True, stdin_open=True, stdout=True, stderr=True)
            print('Started container...', container)
            # Interact with the container using subprocess (for simplicity)
            #  More robust solutions might use docker's exec API for better control.
            process = subprocess.Popen(['docker', 'attach', container.name],
                                       stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
            print('Seems to be ok...')
            # Обработка команд

            # process.stdin.write('help\n'.encode())
            # process.stdin.flush()

            stdio.flush()
            stdio.write('Welcome to custom SSH-server.\n\r')
            stdio.flush()

            while True:
                # stdio.write('{')
                # stdio.flush()

                rlist, wlist, xlist = select.select(
                    [process.stdout, process.stderr, channel],
                    [process.stdin],
                    []
                )

                for fd in rlist:
                    if fd is process.stdout:
                        # output = fd.readline().decode()
                        while True:
                            line: str = fd.readline().decode()
                            line = line.replace('\n', '\n\r')
                            stdio.write(line)
                            stdio.flush()
                            fd.flush()
                            if not line or line == "" or '\x00' in line:
                                break
                    elif fd is process.stderr:
                        while True:
                            line: str = fd.readline().decode()
                            line = line.replace('\n', '\n\r')
                            print('Error:' + line, end='')
                            stdio.write(line)
                            stdio.flush()
                            fd.flush()
                            if not line or line == '' or '\x00' in line:
                                break
                    elif fd is channel:  # Handle input from Paramiko channel
                        try:
                            command = self.command_process(channel)
                            print('Got command: ', command)

                            # if command == 'next' or command == '':
                            #     continue

                            if command == 'exit':
                                self.exit = True
                                break

                            process.stdin.write((command + '\n').encode())
                            process.stdin.flush()
                        except EOFError:
                            print("SSH connection closed.")
                            break
                        except Exception as e:
                            print(f"Error receiving command from channel: {e}")
                            break

                # stdio.write('}')
                # stdio.flush()

                if self.exit:
                    break

            process.stdin.close()
            process.wait()

            container.stop()
            container.remove()

        except docker.errors.BuildError as e:
            print(f"Docker build error: {e}")
        except docker.errors.ContainerError as e:
            print(f"Docker container error: {e}")
        except Exception as e:
            print(f"An error occurred: {e}")

    def command_process(self, channel: Channel):
        channel.send(b'\n\r? ')
        total = b''
        while True:
            c = channel.recv(1024)
            print(c)
            # Backspace
            if c == b'\x7f':
                if len(total) > 0:
                    total = total[:-1]
                    # Сделать бэкспейс, вписать пробел, затем передвинуть курсор
                    channel.send(b'\b \x1b[D')
            # Перенос строки
            elif c == b'\r':
                channel.send(b'\r\n')
                # print('total', total.decode('utf-8'))
                return total.decode('utf-8')
            # Символы
            # ord(c) in range(ord('A'), ord('z') + 1):
            elif len(c) == 1 and c.isascii():
                total += c
                channel.send(c)


if __name__ == '__main__':
    server = SshServer(os.path.expanduser('~/.ssh/id_rsa'))
    server.start("0.0.0.0", int(os.getenv("SSH_PORT", 2222)))