329 lines
12 KiB
Python
329 lines
12 KiB
Python
#!/usr/bin/env python3
|
||
"""
|
||
Proxmox 虚拟机任务控制器
|
||
小主机待机监听 -> 接收命令 -> 启动VM -> 执行任务 -> 汇报进展 -> 关闭VM
|
||
"""
|
||
|
||
import os
|
||
import sys
|
||
import time
|
||
import json
|
||
import logging
|
||
import subprocess
|
||
from pathlib import Path
|
||
from typing import Dict, Any, Optional
|
||
from dataclasses import dataclass
|
||
|
||
import requests
|
||
import paramiko
|
||
from dotenv import load_dotenv
|
||
|
||
# 配置日志
|
||
logging.basicConfig(
|
||
level=logging.INFO,
|
||
format='%(asctime)s - %(levelname)s - %(message)s',
|
||
handlers=[
|
||
logging.FileHandler('/var/log/proxmox_task_controller.log'),
|
||
logging.StreamHandler(sys.stdout)
|
||
]
|
||
)
|
||
logger = logging.getLogger(__name__)
|
||
|
||
@dataclass
|
||
class Config:
|
||
"""配置类"""
|
||
proxmox_host: str
|
||
proxmox_api_token: str
|
||
vm_id: int
|
||
vm_ssh_host: str
|
||
vm_ssh_port: int
|
||
vm_ssh_user: str
|
||
vm_ssh_key_path: str
|
||
command_dir: str
|
||
status_report_url: Optional[str] = None
|
||
|
||
class ConfigLoader:
|
||
"""配置加载器"""
|
||
@staticmethod
|
||
def load() -> Config:
|
||
load_dotenv()
|
||
return Config(
|
||
proxmox_host=os.getenv('PROXMOX_HOST', 'https://localhost:8006'),
|
||
proxmox_api_token=os.getenv('PROXMOX_API_TOKEN', ''),
|
||
vm_id=int(os.getenv('VM_ID', '100')),
|
||
vm_ssh_host=os.getenv('VM_SSH_HOST', 'localhost'),
|
||
vm_ssh_port=int(os.getenv('VM_SSH_PORT', '22')),
|
||
vm_ssh_user=os.getenv('VM_SSH_USER', 'root'),
|
||
vm_ssh_key_path=os.getenv('VM_SSH_KEY_PATH', '~/.ssh/id_rsa'),
|
||
command_dir=os.getenv('COMMAND_DIR', '/var/task-commands'),
|
||
status_report_url=os.getenv('STATUS_REPORT_URL')
|
||
)
|
||
|
||
class ProxmoxAPI:
|
||
"""Proxmox API 客户端"""
|
||
def __init__(self, config: Config):
|
||
self.config = config
|
||
self.base_url = f"{config.proxmox_host}/api2/json"
|
||
self.headers = {
|
||
'Authorization': f'PVEAPIToken={config.proxmox_api_token}'
|
||
}
|
||
|
||
def start_vm(self) -> bool:
|
||
"""启动虚拟机"""
|
||
try:
|
||
url = f"{self.base_url}/nodes/pve/qemu/{self.config.vm_id}/status/start"
|
||
response = requests.post(url, headers=self.headers, verify=False, timeout=30)
|
||
if response.status_code == 200:
|
||
logger.info(f"✅ VM {self.config.vm_id} 启动命令已发送")
|
||
return True
|
||
else:
|
||
logger.error(f"❌ 启动VM失败: {response.status_code} - {response.text}")
|
||
return False
|
||
except Exception as e:
|
||
logger.error(f"❌ 启动VM异常: {e}")
|
||
return False
|
||
|
||
def shutdown_vm(self) -> bool:
|
||
"""关闭虚拟机"""
|
||
try:
|
||
url = f"{self.base_url}/nodes/pve/qemu/{self.config.vm_id}/status/shutdown"
|
||
response = requests.post(url, headers=self.headers, verify=False, timeout=60)
|
||
if response.status_code == 200:
|
||
logger.info(f"✅ VM {self.config.vm_id} 关闭命令已发送")
|
||
return True
|
||
else:
|
||
logger.error(f"❌ 关闭VM失败: {response.status_code} - {response.text}")
|
||
return False
|
||
except Exception as e:
|
||
logger.error(f"❌ 关闭VM异常: {e}")
|
||
return False
|
||
|
||
def get_vm_status(self) -> Dict[str, Any]:
|
||
"""获取虚拟机状态"""
|
||
try:
|
||
url = f"{self.base_url}/nodes/pve/qemu/{self.config.vm_id}/status/current"
|
||
response = requests.get(url, headers=self.headers, verify=False, timeout=10)
|
||
if response.status_code == 200:
|
||
return response.json().get('data', {})
|
||
return {}
|
||
except Exception as e:
|
||
logger.error(f"❌ 获取VM状态异常: {e}")
|
||
return {}
|
||
|
||
class SSHExecutor:
|
||
"""SSH 执行器"""
|
||
def __init__(self, config: Config):
|
||
self.config = config
|
||
|
||
def execute(self, command: str, timeout: int = 300) -> tuple:
|
||
"""在VM中执行命令"""
|
||
try:
|
||
key_path = os.path.expanduser(self.config.vm_ssh_key_path)
|
||
if not os.path.exists(key_path):
|
||
logger.error(f"❌ SSH密钥不存在: {key_path}")
|
||
return False, f"SSH key not found: {key_path}"
|
||
|
||
client = paramiko.SSHClient()
|
||
client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
|
||
client.connect(
|
||
hostname=self.config.vm_ssh_host,
|
||
port=self.config.vm_ssh_port,
|
||
username=self.config.vm_ssh_user,
|
||
key_filename=key_path,
|
||
timeout=30
|
||
)
|
||
|
||
logger.info(f"🔌 SSH连接到VM成功")
|
||
stdin, stdout, stderr = client.exec_command(command, timeout=timeout)
|
||
|
||
exit_code = stdout.channel.recv_exit_status()
|
||
output = stdout.read().decode('utf-8')
|
||
error = stderr.read().decode('utf-8')
|
||
|
||
client.close()
|
||
|
||
if exit_code == 0:
|
||
logger.info(f"✅ 命令执行成功: {command[:50]}...")
|
||
return True, output
|
||
else:
|
||
logger.error(f"❌ 命令执行失败: {command[:50]}... (exit code: {exit_code})")
|
||
return False, error
|
||
|
||
except Exception as e:
|
||
logger.error(f"❌ SSH执行异常: {e}")
|
||
return False, str(e)
|
||
|
||
class CommandProcessor:
|
||
"""命令处理器"""
|
||
def __init__(self, config: Config):
|
||
self.config = config
|
||
self.proxmox = ProxmoxAPI(config)
|
||
self.ssh = SSHExecutor(config)
|
||
self.processed_commands = set()
|
||
|
||
def load_processed_commands(self):
|
||
"""加载已处理的命令记录"""
|
||
record_file = Path(self.config.command_dir) / '.processed_commands'
|
||
if record_file.exists():
|
||
with open(record_file, 'r') as f:
|
||
self.processed_commands = set(json.load(f))
|
||
|
||
def save_processed_commands(self):
|
||
"""保存已处理的命令记录"""
|
||
record_file = Path(self.config.command_dir) / '.processed_commands'
|
||
with open(record_file, 'w') as f:
|
||
json.dump(list(self.processed_commands), f)
|
||
|
||
def report_status(self, status: str, details: str = ""):
|
||
"""汇报状态"""
|
||
if self.config.status_report_url:
|
||
try:
|
||
payload = {
|
||
'timestamp': time.strftime('%Y-%m-%d %H:%M:%S'),
|
||
'status': status,
|
||
'details': details
|
||
}
|
||
requests.post(self.config.status_report_url, json=payload, timeout=10)
|
||
except Exception as e:
|
||
logger.error(f"❌ 状态汇报失败: {e}")
|
||
|
||
logger.info(f"📊 状态更新: {status} - {details}")
|
||
|
||
def wait_for_vm_ready(self, max_wait: int = 120) -> bool:
|
||
"""等待VM就绪"""
|
||
logger.info(f"⏳ 等待VM {self.config.vm_id} 就绪...")
|
||
start_time = time.time()
|
||
|
||
while time.time() - start_time < max_wait:
|
||
status = self.proxmox.get_vm_status()
|
||
if status.get('status') == 'running':
|
||
# 等待SSH可用
|
||
time.sleep(10)
|
||
try:
|
||
self.ssh.execute('echo "VM ready"', timeout=10)
|
||
logger.info("✅ VM已就绪")
|
||
return True
|
||
except:
|
||
logger.info("⏳ VM运行中，等待SSH...")
|
||
time.sleep(5)
|
||
|
||
logger.error("❌ 等待VM就绪超时")
|
||
return False
|
||
|
||
def process_command(self, command_file: Path):
|
||
"""处理单个命令文件"""
|
||
if command_file.name in self.processed_commands:
|
||
return
|
||
|
||
logger.info(f"📋 处理命令文件: {command_file.name}")
|
||
|
||
try:
|
||
with open(command_file, 'r') as f:
|
||
command_data = json.load(f)
|
||
|
||
task_id = command_data.get('task_id', 'unknown')
|
||
command = command_data.get('command', '')
|
||
|
||
self.report_status('starting_vm', f'Task {task_id}')
|
||
|
||
# 1. 启动VM
|
||
if not self.proxmox.start_vm():
|
||
self.report_status('vm_start_failed', f'Task {task_id}')
|
||
return
|
||
|
||
# 2. 等待VM就绪
|
||
if not self.wait_for_vm_ready():
|
||
self.report_status('vm_ready_timeout', f'Task {task_id}')
|
||
self.proxmox.shutdown_vm()
|
||
return
|
||
|
||
self.report_status('executing_task', f'Task {task_id}')
|
||
|
||
# 3. 执行命令
|
||
success, output = self.ssh.execute(command, timeout=300)
|
||
|
||
if success:
|
||
self.report_status('task_completed', f'Task {task_id}')
|
||
logger.info(f"✅ 任务完成: {task_id}")
|
||
else:
|
||
self.report_status('task_failed', f'Task {task_id}: {output}')
|
||
logger.error(f"❌ 任务失败: {task_id}")
|
||
|
||
# 4. 关闭VM
|
||
self.report_status('shutting_down_vm', f'Task {task_id}')
|
||
self.proxmox.shutdown_vm()
|
||
|
||
# 5. 移动命令文件到完成目录
|
||
done_dir = Path(self.config.command_dir) / 'done'
|
||
done_dir.mkdir(exist_ok=True)
|
||
command_file.rename(done_dir / command_file.name)
|
||
|
||
# 6. 记录已处理
|
||
self.processed_commands.add(command_file.name)
|
||
self.save_processed_commands()
|
||
|
||
except Exception as e:
|
||
logger.error(f"❌ 处理命令异常: {e}")
|
||
self.report_status('command_error', str(e))
|
||
|
||
class TaskController:
|
||
"""主任务控制器"""
|
||
def __init__(self):
|
||
self.config = ConfigLoader.load()
|
||
self.processor = CommandProcessor(self.config)
|
||
self.running = True
|
||
|
||
def setup(self):
|
||
"""初始化设置"""
|
||
logger.info("🚀 Proxmox任务控制器启动")
|
||
logger.info(f"📁 监控命令目录: {self.config.command_dir}")
|
||
|
||
# 创建命令目录
|
||
Path(self.config.command_dir).mkdir(parents=True, exist_ok=True)
|
||
|
||
# 加载已处理命令记录
|
||
self.processor.load_processed_commands()
|
||
|
||
# 检查Proxmox连接
|
||
try:
|
||
status = self.processor.proxmox.get_vm_status()
|
||
if status:
|
||
logger.info(f"✅ Proxmox连接正常，VM ID: {self.config.vm_id}")
|
||
else:
|
||
logger.error("❌ Proxmox连接失败")
|
||
except Exception as e:
|
||
logger.error(f"❌ Proxmox连接异常: {e}")
|
||
|
||
def run(self):
|
||
"""主运行循环"""
|
||
self.setup()
|
||
|
||
logger.info("⏳ 进入监听模式，等待命令...")
|
||
|
||
while self.running:
|
||
try:
|
||
command_dir = Path(self.config.command_dir)
|
||
command_files = list(command_dir.glob('*.json'))
|
||
|
||
for cmd_file in command_files:
|
||
if cmd_file.name not in self.processor.processed_commands:
|
||
self.processor.process_command(cmd_file)
|
||
|
||
time.sleep(5)
|
||
|
||
except KeyboardInterrupt:
|
||
logger.info("🛑 收到停止信号")
|
||
self.running = False
|
||
except Exception as e:
|
||
logger.error(f"❌ 主循环异常: {e}")
|
||
time.sleep(10)
|
||
|
||
logger.info("👋 Proxmox任务控制器停止")
|
||
|
||
def main():
|
||
"""主函数"""
|
||
controller = TaskController()
|
||
controller.run()
|
||
|
||
if __name__ == '__main__':
|
||
main() |