spug/spug_api/manage.py

# coding=utf-8
from functools import wraps
from getpass import getpass
import sys
import os

commands = {}


# 工具函数
def check_input_password(password):
    if len(password.strip()) < 6:
        return check_input_password(getpass('密码长度必须大于6位，请重新输入管理员账户密码：'))
    password2 = getpass('再次输入：')
    if password == password2.strip():
        return password
    else:
        return check_input_password(getpass('两次输入密码不一致，请重新输入管理员账户密码：'))


def check_input_username(username):
    if username.strip():
        return username
    else:
        return check_input_username(input('请输入管理员账户登录名：'))


# 注册命令
def registry_command(cmd_str):
    def decorate(func):
        commands[cmd_str] = func

        @wraps(func)
        def wrapper(*args, **kwargs):
            return func(*args, **kwargs)
        return wrapper
    return decorate


@registry_command('init_db')
def init_db():
    from public import db
    from config import BASE_DIR
    import apps.account.models
    import apps.configuration.models
    import apps.deploy.models
    import apps.assets.models
    import apps.schedule.models
    import apps.setting.models

    user_input = input('是否要初始化数据库，该操作会清空所有数据[y|n]？')
    if user_input.strip() == 'y':
        db.drop_all()
        db.create_all()
        with open(os.path.join(BASE_DIR, 'libs', 'sql', 'permissions.sql'), 'r') as f:
            line = f.readline()
            while line:
                if line.startswith('INSERT INTO'):
                    db.engine.execute(line.strip())
                line = f.readline()
        print('数据库已初始化成功！')
        user_input = input('是否需要创建管理员账户[y|n]？')
        if user_input.strip() == 'y':
            create_admin()


@registry_command('create_admin')
def create_admin():
    from apps.account.models import User

    admin = User.query.filter_by(is_supper=True).first()
    if admin:
        user_input = input('已存在管理员账户 <%s>，需要重置密码[y|n]？' % admin.username)
        if user_input.strip() == 'y':
            password = check_input_password(getpass('请输入新的管理员账户密码：'))
            admin.password = password
            admin.token_expired = 0
            admin.save()
            print('重置管理员密码成功！')
    else:
        username = check_input_username(input('请输入管理员账户登录名：'))
        password = check_input_password(getpass('请输入管理员账户密码：'))
        User(username=username, password=password, nickname='管理员', is_supper=True).save()
        print('创建管理员账户成功！')


@registry_command('enable_admin')
def enable_admin():
    from apps.account.models import User

    admin = User.query.filter_by(is_supper=True).first()
    admin.update(is_active=True)
    print('管理员账户状态已修改为启用！')


def print_usage():
    print('''
usage: %s <command>

command:
    init_db         初始化数据库
    create_admin    创建管理员账户
    enable_admin    启用管理员账户，用于登录失败次数过多账户被禁用时使用
    ''' % sys.argv[0])


if __name__ == '__main__':
    if len(sys.argv) == 1:
        print_usage()
        sys.exit(1)
    cmd = sys.argv.pop(0)
    arg1 = sys.argv.pop(0)
    r_func = commands.get(arg1)
    if callable(r_func):
        r_func(*sys.argv)
    else:
        print('遇到了不可能会出现的错误！')