#!/usr/bin/env python3
import asyncio
import asyncssh
import os
import sys
from io import BytesIO

# Класс сессии: обрабатывает подключение и ввод данных
class MySSHServerSession(asyncssh.SSHServerSession):
    def connection_made(self, chan):
        self.chan = chan

    def session_started(self):
        self.chan.write("Добро пожаловать в SSH-сервис на Hugging Face Spaces!\n")
        self.chan.write("Введите команду (для выхода наберите 'exit'):\n")

    def data_received(self, data, datatype):
        # Если пользователь ввёл 'exit', закрываем сессию
        if data.strip() == 'exit':
            self.chan.write("До свидания!\n")
            self.chan.exit(0)
        else:
            # Эхо-ответ: выводим, что ввёл пользователь
            self.chan.write("Вы ввели: " + data)

    def connection_lost(self, exc):
        pass

# Класс SSH-сервера
class MySSHServer(asyncssh.SSHServer):
    def session_requested(self):
        return MySSHServerSession()

    async def begin_auth(self, username):
        # Простой способ проверки: разрешаем любые логины
        return True

# Генерация временного SSH-ключа в памяти
def generate_temp_ssh_key():
    private_key = asyncssh.generate_private_key('ssh-rsa')
    private_key_str = private_key.private_key_data.decode('utf-8')  # Используем правильное свойство
    return private_key_str

# Функция для запуска сервера
async def start_ssh_server():
    # Читаем порт из переменной окружения PORT (по умолчанию 2222)
    port = int(os.environ.get('PORT', 2222))

    # Генерация временного ключа в памяти
    private_key_str = generate_temp_ssh_key()

    try:
        # Используем приватный ключ из памяти
        private_key = asyncssh.import_private_key(private_key_str)

        server = await asyncssh.create_server(
            MySSHServer, '', port,
            server_host_keys=[private_key]
        )
    except (OSError, asyncssh.Error) as exc:
        sys.exit('Ошибка при запуске сервера: ' + str(exc))
    
    print('SSH-сервер запущен на порту', port)
    await server.wait_closed()

def main():
    try:
        asyncio.run(start_ssh_server())
    except (OSError, asyncssh.Error) as exc:
        sys.exit('Ошибка: ' + str(exc))

if __name__ == '__main__':
    main()