diviner/diviner.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
import os
import sys
import subprocess
from demeter.core import *

def init():
    #import nltk
    #nltk.download('punkt', quiet=True)
    #nltk.download('averaged_perceptron_tagger', quiet=True)

    model = Demeter.model('manage_admin')
    model.id = 1
    info = model.select(type='fetchone')
    if not info:
        model.role_id = 1
        model.username = 'admin'
        model.mobile = '15810090845'
        model.password = '123456'
        model.insert()

    model = Demeter.model('manage_role')
    model.id = 1
    info = model.select(type='fetchone')
    if not info:
        model.name = u'管理员'
        model.insert()

def get_service_name(script_name):
    return 'diviner_' + script_name.replace(".py", "")

def install(script_name):
    init()
    bin_path = '/usr/local/bin'
    target_path = os.path.abspath(script_name)
    link_path = os.path.join(bin_path, os.path.splitext(os.path.basename(script_name))[0])

    Demeter.echo(f"安装脚本到 {link_path}")
    if os.path.exists(link_path):
        subprocess.run(['sudo', 'rm', '-f', link_path], check=True)
    subprocess.run(['sudo', 'ln', '-s', target_path, link_path], check=True)
    subprocess.run(['sudo', 'chmod', '+x', target_path], check=True)
    Demeter.echo(f"安装完成，可以直接用命令 `{os.path.basename(script_name).replace('.py','')} start` 启动")

def start(script_name, desc):
    template = """[Unit]
Description={desc}
After=network.target

[Service]
ExecStart={python_path} {workdir}/{script}
WorkingDirectory={workdir}
Restart=always
RestartSec=3
User={user}
Environment=PYTHONUNBUFFERED=1

[Install]
WantedBy=multi-user.target
"""
    workdir = os.getcwd()
    python_path = subprocess.getoutput("which python3")
    user = subprocess.getoutput("whoami")
    service_name = get_service_name(script_name)
    service_file = f"/etc/systemd/system/{service_name}.service"
    
    content = template.format(
        desc=desc,
        python_path=python_path,
        workdir=workdir,
        script=script_name,
        user=user,
    )

    # 写入临时 service 文件
    with open(f"/tmp/{service_name}.service", "w") as f:
        f.write(content)

    # 移动到 systemd 目录（需 root 权限）
    subprocess.run(["sudo", "mv", f"/tmp/{service_name}.service", service_file], check=True)

    # 重新加载 systemd，启用并启动服务
    subprocess.run(["sudo", "systemctl", "daemon-reload"], check=True)
    subprocess.run(["sudo", "systemctl", "enable", service_name], check=True)
    subprocess.run(["sudo", "systemctl", "restart", service_name], check=True)

    Demeter.echo(f"已安装并启动服务: {service_name}")

def stop(script_name):
    service_name = get_service_name(script_name)
    Demeter.echo(f"停止服务: {service_name}")
    subprocess.run(["sudo", "systemctl", "stop", service_name], check=True)
    subprocess.run(["sudo", "systemctl", "disable", service_name], check=True)

def restart(script_name):
    service_name = get_service_name(script_name)
    Demeter.echo(f"重启服务: {service_name}")
    subprocess.run(["sudo", "systemctl", "restart", service_name], check=True)

def main():
    services = [
        ("front.py", "Diviner Front Service"),
        ("admin.py", "Diviner Admin Service"),
        ("cron.py", "Diviner Cron Service"),
    ]

    action = "start"
    if len(sys.argv) > 1:
        action = sys.argv[1].lower()

    if action == "install":
        current_script = os.path.basename(sys.argv[0])
        if not os.path.exists(current_script):
            Demeter.echo(f"当前脚本文件不存在: {current_script}")
            sys.exit(1)
        install(current_script)
    else:
        valid_services = [(s, d) for s, d in services if os.path.exists(s)]
        for s, d in services:
            if not os.path.exists(s):
                Demeter.echo(f"脚本不存在: {s}")

        if action == "start":
            for script, desc in valid_services:
                start(script, desc)
        elif action == "stop":
            for script, _ in valid_services:
                stop(script)
        elif action == "restart":
            for script, _ in valid_services:
                restart(script)
        else:
            Demeter.echo(f"未知动作: {action}")
            sys.exit(1)

if __name__ == "__main__":
    main()