Update vps_monitor.py

vps_monitor.py

@@ -6,21 +6,32 @@ import sys
 from flask import Flask, jsonify, render_template_string
 from threading import Thread
 import logging
+from colorama import init, Fore, Back, Style
+
+init(autoreset=True)  # 初始化 colorama
 
 app = Flask(__name__)
 
 vps_status = {}
 
 # 设置日志
-logging.basicConfig(
-    level=logging.INFO,
-    format='%(asctime)s - %(levelname)s - %(message)s',
-    handlers=[
-        logging.StreamHandler(sys.stdout),
-        logging.StreamHandler(sys.stderr)
-    ]
-)
+class ColoredFormatter(logging.Formatter):
+    def format(self, record):
+        if record.levelno == logging.INFO:
+            record.msg = f"{Fore.GREEN}{record.msg}{Style.RESET_ALL}"
+        elif record.levelno == logging.WARNING:
+            record.msg = f"{Fore.YELLOW}{record.msg}{Style.RESET_ALL}"
+        elif record.levelno == logging.ERROR:
+            record.msg = f"{Fore.RED}{record.msg}{Style.RESET_ALL}"
+        return super().format(record)
+
+formatter = ColoredFormatter('%(asctime)s - %(levelname)s - %(message)s')
+handler = logging.StreamHandler(sys.stdout)
+handler.setFormatter(formatter)
+
 logger = logging.getLogger()
+logger.setLevel(logging.INFO)
+logger.addHandler(handler)
 
 def get_vps_configs():
     configs = []
@@ -52,18 +63,13 @@ def get_vps_configs():
             }
             configs.append(config)
         
-        # 添加调试输出
         logger.info(f"Config {index}: {hostname}, {username}, {script_paths}")
-        print(f"Config {index}: {hostname}, {username}, {script_paths}")
-        sys.stdout.flush()
         
         index += 1
     return configs
 
 def check_and_run_script(config):
     logger.info(f"Checking VPS {config['index']}: {config['hostname']} - {config['script_path']}")
-    print(f"Checking VPS {config['index']}: {config['hostname']} - {config['script_path']}")
-    sys.stdout.flush()
     client = None
     try:
         client = paramiko.SSHClient()
@@ -80,24 +86,25 @@ def check_and_run_script(config):
         script_path = config['script_path']
         script_name = os.path.basename(script_path)
         
-        # 检查脚本是否正在运行 (使用 ps 命令，适用于FreeBSD)
         check_command = f"ps aux | grep {script_name} | grep -v grep"
         stdin, stdout, stderr = client.exec_command(check_command)
         if stdout.read():
-            status = "Running"
+            status = f"{Fore.GREEN}Running{Style.RESET_ALL}"
             logger.info(f"Script {script_name} is running on {config['hostname']}")
         else:
-            logger.info(f"Script {script_name} not running on {config['hostname']}. Executing restart script.")
-            restart_command = f"/bin/sh {script_path}"  # 使用完整路径的sh
+            logger.warning(f"Script {script_name} not running on {config['hostname']}. Executing restart script.")
+            restart_command = f"/bin/sh {script_path}"
             stdin, stdout, stderr = client.exec_command(restart_command)
             exit_status = stdout.channel.recv_exit_status()
             
             if exit_status == 0:
-                status = "Restarted"
+                status = f"{Fore.YELLOW}Restarted{Style.RESET_ALL}"
                 logger.info(f"Restart script {script_name} executed successfully on {config['hostname']}")
+                output = stdout.read().decode('utf-8')
+                logger.info(f"Restart output: {output}")
             else:
                 error_output = stderr.read().decode('utf-8')
-                status = f"Error: {error_output}"
+                status = f"{Fore.RED}Error{Style.RESET_ALL}"
                 logger.error(f"Error executing restart script {script_name} on {config['hostname']}: {error_output}")
         
         key = f"{config['hostname']}:{script_name}"
@@ -114,7 +121,7 @@ def check_and_run_script(config):
         key = f"{config['hostname']}:{script_name}"
         vps_status[key] = {
             'index': config['index'],
-            'status': f"Error: {str(e)}",
+            'status': f"{Fore.RED}Error: {str(e)}{Style.RESET_ALL}",
             'last_check': time.strftime('%Y-%m-%d %H:%M:%S'),
             'username': config['username'],
             'script_name': script_name
@@ -124,20 +131,15 @@ def check_and_run_script(config):
             client.close()
             logger.info(f"SSH connection closed for VPS {config['index']}: {config['hostname']}")
         logger.info(f"Finished checking VPS {config['index']}: {config['hostname']} - {script_name}")
-        print(f"Finished checking VPS {config['index']}: {config['hostname']} - {script_name}")
-        sys.stdout.flush()
 
 def check_all_vps():
-    logger.info("Starting VPS check")
-    print("Starting VPS check")
-    sys.stdout.flush()
+    logger.info(f"{Back.BLUE}{Fore.WHITE}Starting VPS check{Style.RESET_ALL}")
     vps_configs = get_vps_configs()
     for config in vps_configs:
         check_and_run_script(config)
-    logger.info("Finished VPS check")
+    logger.info(f"{Back.BLUE}{Fore.WHITE}Finished VPS check{Style.RESET_ALL}")
     for key, status in vps_status.items():
-        print(f"VPS {status['index']} - {key}: Status: {status['status']}, Username: {status['username']}, Script: {status['script_name']}")
-    sys.stdout.flush()
+        logger.info(f"VPS {status['index']} - {key}: Status: {status['status']}, Username: {status['username']}, Script: {status['script_name']}")
 
 @app.route('/')
 def index():
@@ -184,37 +186,24 @@ def main():
     global start_time
     start_time = time.time()
     
-    print("===== VPS monitoring script is starting =====")
-    sys.stdout.flush()
-    logger.info("===== VPS monitoring script started =====")
+    logger.info("===== VPS monitoring script is starting =====")
     
     flask_thread = Thread(target=run_flask)
     flask_thread.start()
     logger.info("Flask server started in background")
-    print("Flask server started in background")
-    sys.stdout.flush()
 
     vps_configs = get_vps_configs()
     logger.info(f"Found {len(vps_configs)} VPS configurations")
-    print(f"Found {len(vps_configs)} VPS configurations")
-    sys.stdout.flush()
     for config in vps_configs:
         logger.info(f"VPS configured: {config['hostname']} - {config['script_path']}")
-        print(f"VPS configured: {config['hostname']} - {config['script_path']}")
-        sys.stdout.flush()
 
     logger.info("Running initial VPS check")
-    print("Running initial VPS check")
-    sys.stdout.flush()
     check_all_vps()
 
     schedule.every(15).minutes.do(check_all_vps)
     logger.info("Scheduled VPS check every 15 minutes")
-    print("Scheduled VPS check every 15 minutes")
-    sys.stdout.flush()
     
-    print("===== VPS monitoring script is running =====")
-    sys.stdout.flush()
+    logger.info("===== VPS monitoring script is running =====")
     
     heartbeat_count = 0
     while True:
@@ -222,8 +211,7 @@ def main():
         time.sleep(60)
         heartbeat_count += 1
         if heartbeat_count % 5 == 0:  # 每5分钟输出一次心跳信息
-            print(f"Heartbeat: Script is still running. Uptime: {heartbeat_count} minutes")
-            sys.stdout.flush()
+            logger.info(f"Heartbeat: Script is still running. Uptime: {heartbeat_count} minutes")
 
 if __name__ == "__main__":
     main()