pve9-repo-fix/RDP/rdp-toggle-api.py

#!/usr/bin/env python3
"""
RDP Toggle API Server
Control RDP/Shell display mode via REST API
"""

from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
import subprocess
import os
import json
from typing import Optional
from datetime import datetime
import uvicorn

app = FastAPI(title="RDP Toggle API", version="1.0.0")

# CORS 설정 (외부에서 접근 가능)
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# 상태 저장 파일
STATE_FILE = "/var/lib/rdp-toggle/state.json"
CONFIG_FILE = "/var/lib/rdp-toggle/config.json"

# 기본 설정
DEFAULT_CONFIG = {
    "rdp_server": "192.168.0.229",
    "rdp_username": "0bin",
    "rdp_password": "trajet6640",
    "local_user": "rdpuser"
}

class ToggleRequest(BaseModel):
    mode: str  # "rdp" or "shell"
    
class ConfigUpdate(BaseModel):
    rdp_server: Optional[str] = None
    rdp_username: Optional[str] = None
    rdp_password: Optional[str] = None
    local_user: Optional[str] = None

class StatusResponse(BaseModel):
    current_mode: str
    rdp_active: bool
    last_changed: str
    config: dict

def ensure_directories():
    """필요한 디렉토리 생성"""
    os.makedirs("/var/lib/rdp-toggle", exist_ok=True)

def load_state():
    """현재 상태 로드"""
    if os.path.exists(STATE_FILE):
        with open(STATE_FILE, 'r') as f:
            return json.load(f)
    return {
        "current_mode": "shell",
        "rdp_active": False,
        "last_changed": datetime.now().isoformat()
    }

def save_state(state):
    """상태 저장"""
    ensure_directories()
    with open(STATE_FILE, 'w') as f:
        json.dump(state, f)

def load_config():
    """설정 로드"""
    if os.path.exists(CONFIG_FILE):
        with open(CONFIG_FILE, 'r') as f:
            return json.load(f)
    return DEFAULT_CONFIG

def save_config(config):
    """설정 저장"""
    ensure_directories()
    with open(CONFIG_FILE, 'w') as f:
        json.dump(config, f)

def enable_rdp():
    """RDP 모드 활성화"""
    config = load_config()
    
    # 자동 로그인 설정
    subprocess.run([
        "mkdir", "-p", "/etc/systemd/system/getty@tty1.service.d"
    ], check=False)
    
    override_content = f"""[Service]
ExecStart=
ExecStart=-/sbin/agetty --autologin {config['local_user']} --noclear %I $TERM
Type=idle"""
    
    with open("/etc/systemd/system/getty@tty1.service.d/override.conf", "w") as f:
        f.write(override_content)
    
    # X 자동 시작 스크립트
    bash_profile = f"""if [[ -z $DISPLAY ]] && [[ $(tty) == /dev/tty1 ]]; then
    startx
    logout
fi"""
    
    user_home = f"/home/{config['local_user']}"
    with open(f"{user_home}/.bash_profile", "w") as f:
        f.write(bash_profile)
    
    # .xinitrc 업데이트
    xinitrc_content = f"""#!/bin/bash
xset -dpms
xset s off
xset s noblank
unclutter -idle 0.1 -root &
openbox-session &
sleep 2
xfreerdp3 /v:{config['rdp_server']} /u:{config['rdp_username']} /p:"{config['rdp_password']}" +f /cert:ignore +dynamic-resolution /sound:sys:alsa +clipboard
pkill -SIGTERM Xorg"""
    
    with open(f"{user_home}/.xinitrc", "w") as f:
        f.write(xinitrc_content)
    
    subprocess.run(["chmod", "+x", f"{user_home}/.xinitrc"])
    subprocess.run(["chown", f"{config['local_user']}:{config['local_user']}", 
                   f"{user_home}/.bash_profile", f"{user_home}/.xinitrc"])
    
    # systemd 리로드 및 getty 재시작
    subprocess.run(["systemctl", "daemon-reload"])
    subprocess.run(["systemctl", "restart", "getty@tty1.service"])
    
    return True

def disable_rdp():
    """Shell 모드로 전환 (RDP 비활성화)"""
    config = load_config()
    
    # RDP 프로세스 종료
    subprocess.run(["pkill", "-u", config['local_user'], "xfreerdp3"], check=False)
    subprocess.run(["pkill", "-u", config['local_user'], "-f", "xinit|Xorg|openbox"], check=False)
    
    # 자동 로그인 설정 제거
    subprocess.run(["rm", "-f", "/etc/systemd/system/getty@tty1.service.d/override.conf"], check=False)
    
    # 자동 시작 스크립트 제거
    user_home = f"/home/{config['local_user']}"
    subprocess.run(["rm", "-f", f"{user_home}/.bash_profile"], check=False)
    
    # systemd 리로드 및 getty 재시작
    subprocess.run(["systemctl", "daemon-reload"])
    subprocess.run(["systemctl", "restart", "getty@tty1.service"])
    
    # TTY1으로 전환
    subprocess.run(["chvt", "1"], check=False)
    
    return True

@app.get("/")
async def root():
    """API 정보"""
    return {
        "name": "RDP Toggle API",
        "version": "1.0.0",
        "endpoints": {
            "GET /status": "현재 상태 확인",
            "POST /toggle": "모드 전환 (rdp/shell)",
            "GET /config": "현재 설정 확인",
            "PUT /config": "설정 업데이트"
        }
    }

@app.get("/status", response_model=StatusResponse)
async def get_status():
    """현재 상태 반환"""
    state = load_state()
    config = load_config()
    
    # 실제 프로세스 확인
    try:
        result = subprocess.run(
            ["pgrep", "-f", "xfreerdp3"],
            capture_output=True,
            text=True
        )
        rdp_running = result.returncode == 0
        state["rdp_active"] = rdp_running
        state["current_mode"] = "rdp" if rdp_running else "shell"
    except:
        pass
    
    return StatusResponse(
        current_mode=state["current_mode"],
        rdp_active=state["rdp_active"],
        last_changed=state["last_changed"],
        config=config
    )

@app.post("/toggle")
async def toggle_mode(request: ToggleRequest):
    """모드 전환"""
    if request.mode not in ["rdp", "shell"]:
        raise HTTPException(status_code=400, detail="Mode must be 'rdp' or 'shell'")
    
    state = load_state()
    
    try:
        if request.mode == "rdp":
            success = enable_rdp()
            if success:
                state["current_mode"] = "rdp"
                state["rdp_active"] = True
        else:  # shell
            success = disable_rdp()
            if success:
                state["current_mode"] = "shell"
                state["rdp_active"] = False
        
        state["last_changed"] = datetime.now().isoformat()
        save_state(state)
        
        return {
            "status": "success",
            "mode": request.mode,
            "message": f"Switched to {request.mode} mode"
        }
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.get("/config")
async def get_config():
    """현재 설정 반환"""
    return load_config()

@app.put("/config")
async def update_config(update: ConfigUpdate):
    """설정 업데이트"""
    config = load_config()
    
    if update.rdp_server:
        config["rdp_server"] = update.rdp_server
    if update.rdp_username:
        config["rdp_username"] = update.rdp_username
    if update.rdp_password is not None:
        config["rdp_password"] = update.rdp_password
    if update.local_user:
        config["local_user"] = update.local_user
    
    save_config(config)
    
    # 현재 RDP 모드인 경우 재시작
    state = load_state()
    if state["rdp_active"]:
        disable_rdp()
        enable_rdp()
    
    return {
        "status": "success",
        "config": config
    }

if __name__ == "__main__":
    # Root 권한 확인
    if os.geteuid() != 0:
        print("This script must be run as root")
        exit(1)
    
    ensure_directories()
    
    # 서버 시작
    uvicorn.run(app, host="0.0.0.0", port=8090)