ssh_manager.py

import asyncio
import asyncssh
from db import save_session, delete_session, update_timestamp
from config import SESSION_TIMEOUT

active_sessions = {}


async def ssh_reader(uid):
    sess = active_sessions.get(uid)
    if not sess:
        return

    proc = sess["proc"]
    queue = sess["queue"]

    try:
        while True:
            line = await proc.stdout.readline()
            if not line:
                break
            await queue.put(line)
    except Exception:
        pass


async def start_session(uid, ip, port, user, pw):
    conn = await asyncssh.connect(
        ip,
        port=port,
        username=user,
        password=pw,
        known_hosts=None
    )

    proc = await conn.create_process()
    queue = asyncio.Queue()

    active_sessions[uid] = {
        "conn": conn,
        "proc": proc,
        "queue": queue,
        "task": None
    }

    task = asyncio.create_task(ssh_reader(uid))
    active_sessions[uid]["task"] = task

    await save_session(uid, ip, port, user, pw)


async def close_session(uid):
    sess = active_sessions.get(uid)
    if not sess:
        return

    try:
        sess["proc"].stdin.write("exit\n")
    except:
        pass

    try:
        sess["conn"].close()
    except:
        pass

    try:
        if sess["task"]:
            sess["task"].cancel()
    except:
        pass

    active_sessions.pop(uid, None)
    await delete_session(uid)


async def send_command(uid, text):
    sess = active_sessions.get(uid)
    if not sess:
        return None

    sess["proc"].stdin.write(text + "\n")
    await update_timestamp(uid)

    output = ""
    for _ in range(25):
        try:
            line = sess["queue"].get_nowait()
            output += line
        except asyncio.QueueEmpty:
            break

    return output or " "