服务

Python工具

# -*- coding: utf8 -*-
"""
Centos7 下注册为自启动服务机制
支持start, restart, stop命令：systemctl start|restart|stop my_test_server
1. 在工作路径下生成sh
2. 把sh注册为服务
3. 通过sh控制python
"""

import os
import sys
import time
import subprocess

version = '0.0.2'

server_content = '''[Unit]
Description=%(description)s
After=network.target sshd.service

[Service]
Type=simple
ExecStart=%(sh)s start &
ExecReload=%(sh)s restart &
ExecStop=%(sh)s stop
Restart=always
RestartSec=10s

[Install]
WantedBy=multi-user.target
'''


def on_cmd(cmd):
    p = subprocess.Popen(cmd, stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.PIPE, shell=True)
    d = p.stdout.read()
    if p.wait() != 0:
        # logging.error("cmd error: %s, returncode: %d", cmd, p.returncode)
        return None
    return d


def get_py_cmd(args):
    cmd = ''
    for arg in args:
        if arg:
            cmd += ' ' + arg
    return cmd


def install_sh(name, fn, work_path, args, python_args=None):
    with open(os.path.join(os.path.dirname(__file__), 'service.sh'), 'r') as f:
        content = f.read()
        content = content % {'cmd': get_py_cmd(['python', python_args, fn, args])}
        fn_sh = os.path.join(work_path, name + '.service.sh')
        with open(fn_sh, 'w') as sf:
            sf.write(content)
            return fn_sh


def install(name, fn, args=None, work_path=None, description=None, python_args=None):
    # fn是相对路径
    # dir是工作路径: 默认fn所在位置
    if not work_path:
        work_path = os.path.dirname(fn)
    work_path = os.path.abspath(work_path)
    fn_sh = install_sh(name, fn, work_path, args, python_args)
    if not fn_sh:
        return False
    with open('/usr/lib/systemd/system/%s.service' % name, 'w') as f:
        d = {
            'name': name,
            'sh': fn_sh,
            'description': description
        }
        f.write(server_content % d)
        on_cmd("chmod +x %s" % fn_sh)
        return True
    return False


def uninstall(name):
    fn_service = '/usr/lib/systemd/system/%s.service' % name
    on_cmd("rm -rf %s" % fn_service)


def check_exist(name):
    d = on_cmd("ls /usr/lib/systemd/system/")
    if d.find("%s.service" % name) != -1:
        return True
    return False


def test():
    name = 'my_test_server'
    if len(sys.argv) == 2:
        if sys.argv[1] == 'install':
            if not check_exist(name):
                fn = __file__
                fn = os.path.split(fn)
                work_path = os.path.dirname(fn[0])  # 上级路径
                fn = os.path.basename(fn[0]) + '.' + os.path.splitext(fn[1])[0]  # 模块名称
                if install(name, fn, args="test", work_path=work_path,
                           description="my_startup test", python_args='-m'):
                    print 'install success'
                else:
                    print 'install fail'
            else:
                print 'is install'
        elif sys.argv[1] == 'uninstall':
            if check_exist(name):
                on_cmd("rm -rf /usr/lib/systemd/system/%s.service" % name)
                print 'uninstall success'
            else:
                print 'is uninstall'
        elif sys.argv[1] == 'test':
            while 1:
                print 'test'
                time.sleep(10)
    else:
        print "arg: install/uninstall"


if __name__ == "__main__":
    test()