core/ssh_agent.py - master - HookDev-Arch/ServerMonitor

core/ssh_agent.py
66 строк2 KB
HookDev-Arch
upload files
08 окт 2025, 14:53
08 окт 2025, 14:53f2498cb
import time
import requests
import paramiko
from datetime import datetime
from pathlib import Path
from .utils import LOGS_DIR

AGENT_CODE = """#!/usr/bin/env python3
import psutil, json
from datetime import datetime
from fastapi import FastAPI
app = FastAPI()
@app.get('/api/stats')
def stats():
    mem = psutil.virtual_memory()
    return {"timestamp": datetime.now().isoformat(), "memory": {"used": mem.used, "total": mem.total}}
if __name__ == '__main__':
    import uvicorn; uvicorn.run(app, host='0.0.0.0', port=3333)
"""

def ssh_connect(host, port, user, password, timeout=8):
    client = paramiko.SSHClient()
    client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
    client.connect(hostname=host, port=port, username=user, password=password, timeout=timeout)
    return client

def ssh_exec(client, cmd):
    stdin, stdout, stderr = client.exec_command(cmd)
    return stdout.read().decode(), stderr.read().decode()

def deploy_agent(server):
    host, port, user, password = server["host"], server["port"], server["user"], server["password"]
    api_port = server.get("api_port", 3333)
    path = server.get("agent_path", "~/agent.py")

    client = ssh_connect(host, port, user, password)
    sftp = client.open_sftp()
    sftp_file = sftp.file(path.replace("~", "/root"), "w")
    sftp_file.write(AGENT_CODE)
    sftp_file.close()
    sftp.close()

    ssh_exec(client, f"nohup python3 {path} > /dev/null 2>&1 &")
    time.sleep(2)
    try:
        r = requests.get(f"http://{host}:{api_port}/api/stats", timeout=3)
        if r.ok:
            return True
    except Exception:
        return False
    finally:
        client.close()
    return False

def fetch_remote_stats(server):
    host = server["host"]
    port = server.get("api_port", 3333)
    url = f"http://{host}:{port}/api/stats"
    try:
        r = requests.get(url, timeout=3)
        data = r.json() if r.ok else {"error": f"HTTP {r.status_code}"}
    except Exception as e:
        data = {"error": str(e)}
    with (LOGS_DIR / f"{server['name']}.jsonl").open("a", encoding="utf-8") as f:
        f.write(json.dumps({"time": datetime.now().isoformat(), "data": data}, ensure_ascii=False) + "\n")
    return data