#!/usr/bin/env python3
# -*- coding: utf-8 -*-
import sys
import os
import csv
import re

# ====================== 基础配置 ======================
IP_FILE = "ip.txt"
SSH_USER = "root"
SSH_PASS = "hp9000"
SSH_TIMEOUT = 3
# ======================================================

# ===================== 自动安装依赖 ====================
def auto_install_python_and_paramiko():
    print("=> 检查环境是否满足...")
    ret = os.system("which python3 >/dev/null 2>&1")
    if ret != 0:
        print("=> 未找到 python3，开始自动安装...")
        # 自动识别系统
        if os.path.exists("/etc/redhat-release"):
            os.system("yum install -y python3 python3-pip")
        elif os.path.exists("/etc/debian_version"):
            os.system("apt update && apt install -y python3 python3-pip")
        else:
            print("不支持的系统")
            sys.exit(1)

    # 安装 paramiko
    print("=> 安装 paramiko SSH 库...")
    os.system("python3 -m pip install paramiko --quiet")

# 先安装依赖
try:
    import paramiko
except ImportError:
    auto_install_python_and_paramiko()
    import paramiko
# =================================================================

def print_help():
    help_text = """
用法：python3 check_proc.py 进程名称
示例：
    python3 check_proc.py trap
    python3 check_proc.py node_exporter
    python3 check_proc.py nginx
"""
    print(help_text)
    sys.exit(1)

def clean_ip_file(filepath):
    clean_ips = []
    with open(filepath, "r", encoding="utf-8") as f:
        for line in f:
            line = line.replace("\r", "").strip()
            if line:
                clean_ips.append(line)
    return clean_ips

def judge_ip_type(ip):
    if ":" in ip:
        return "IPv6"
    return "IPv4"

def ssh_check_host(ip, proc_name, timeout):
    client = paramiko.SSHClient()
    client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
    try:
        client.connect(
            hostname=ip,
            username=SSH_USER,
            password=SSH_PASS,
            timeout=timeout,
            banner_timeout=timeout
        )
        remote_cmd = f'''
hostname
cat /etc/os-release | grep PRETTY_NAME | cut -d'"' -f2 | head -n1
ps -ef | grep -v grep | grep -i '{proc_name}' | wc -l
        '''.strip()

        stdin, stdout, stderr = client.exec_command(remote_cmd)
        output = stdout.read().decode("utf-8", errors="ignore").strip()
        client.close()

        if not output:
            return {"hn": "无输出", "os": "---", "status": "未知"}

        lines = [x.strip() for x in output.splitlines() if x.strip()]
        hn = lines[0] if len(lines) >= 1 else "获取失败"
        os_ver = lines[1] if len(lines) >= 2 else "获取失败"
        proc_count = lines[2] if len(lines) >= 3 else "0"

        try:
            cnt = int(proc_count)
        except:
            cnt = 0

        status = "已运行" if cnt >= 1 else "未运行"
        return {"hn": hn, "os": os_ver, "status": status}

    except paramiko.AuthenticationException:
        return {"hn": "密码错误", "os": "---", "status": "认证失败"}
    except:
        return {"hn": "SSH超时", "os": "---", "status": "超时"}

def main():
    if len(sys.argv) != 2:
        print_help()
    target_proc = sys.argv[1].strip()
    csv_name = f"进程检查结果_{target_proc}.csv"

    ip_list = clean_ip_file(IP_FILE)
    total = len(ip_list)

    print(f"===== 批量进程检查：{target_proc} =====")
    print(f"总IP数：{total}\n")

    with open(csv_name, "w", encoding="utf-8-sig", newline="") as f:
        w = csv.writer(f)
        w.writerow(["IP地址", "IP类型", "主机名", "系统版本", "进程状态"])

        for ip in ip_list:
            typ = judge_ip_type(ip)
            print(f"→ {ip}", end=" ")
            res = ssh_check_host(ip, target_proc, SSH_TIMEOUT)
            print(f"[{res['status']}]")
            w.writerow([ip, typ, res["hn"], res["os"], res["status"]])

    print("\n===== 执行完成 =====")
    print(f"报告：{csv_name}")

if __name__ == "__main__":
    main()