create a user for yapi

README.markdown

create a user for yapi, e.g.:

$ python3 create_user.py --mongo_host localhost --mongo_port 27017 --username haha --useremail [email protected]
INFO: generate default password: Uck33M42Za
INFO: create user successfully, id=29

note:

• install python lib pymongo

    pip3 install pymongo
  

• param username is optional. if username is not specified, the script will abstract username from param useremail.

create_user.py

import re
import time
import random
import string
import hashlib
import argparse
import pymongo

REGEXP_EMAIL = r'^(<\w[\s\w]+>\s)?(\w+[\w+.]*@\w+.(org|com)$)'
PASSWORD_SALT_LENGTH = 11
PASSWROD_PLAIN_TEXT_LENGTH = 10
DEFAULT_USER_ROLE = 'member'
DEFAULT_REGISTER_TYPE = 'site'
MONGO_SERVER_SELECTION_TIMEOUT_IN_MS = 10

def generate_random_str(str_length):
    chars = string.ascii_letters + string.digits
    return ''.join([random.choice(chars) for i in range(str_length)])

def generate_plain_password(password_length):
    return generate_random_str(password_length)

def generate_salt(salt_length):
    return generate_random_str(salt_length)

def check_email_existed(mongo_model_user, useremail):
    existed_user = mongo_model_user.count_documents({'email': useremail})
    if existed_user == 0:
        return False
    else:
        return True

def update_identitycounters(mongo_db_yapi, step):
    model_identitycounters = mongo_db_yapi['identitycounters']
    existed_identitycounter_user = model_identitycounters.count_documents({'model':'user'})
    if existed_identitycounter_user == 0:
        print('please init yapi server with yapi command tool, skip')
        exit(1)

    identitycounter_user = model_identitycounters.find_one({'model':'user'})
    current_count = identitycounter_user['count']
    new_count = current_count + step
    model_identitycounters.update_one({'model':'user'}, {'$set': {'count': new_count}})
    return new_count

def create_user(mongo_host, mongo_port, username, useremail):
    try:
        client = pymongo.MongoClient(host=mongo_host, port=mongo_port, serverSelectionTimeoutMS=MONGO_SERVER_SELECTION_TIMEOUT_IN_MS)
        client.server_info()
    except pymongo.errors.ServerSelectionTimeoutError:
        print('ERROR: connect to mongo failed')
        exit(1)

    db_all = client.list_database_names()
    db_yapi_existed = False
    if 'yapi' in db_all:
        db_yapi_existed = True
    if not db_yapi_existed:
        print('ERROR: mongo db yapi not existed')
        exit(1)

    db_yapi = client['yapi']
    model_user = db_yapi['user']

    if (check_email_existed(model_user, useremail)):
        print('ERROR: email existed, skip')
        exit(1)

    current_time_in_seconds = int(time.time())

    new_user_password_salt = generate_salt(PASSWORD_SALT_LENGTH)
    new_user_password_plan_text = generate_plain_password(PASSWROD_PLAIN_TEXT_LENGTH)
    print('INFO: generate default password: %s' % new_user_password_plan_text)
    new_user_password_encrypted = hashlib.sha1((new_user_password_plan_text +  hashlib.sha1(new_user_password_salt.encode('utf-8')).hexdigest()).encode('utf-8')).hexdigest()
    new_user_id = update_identitycounters(db_yapi, 1)
    
    new_user = {}
    new_user['_id'] = new_user_id
    new_user['username'] = username
    new_user['email'] = useremail
    new_user['passsalt'] = new_user_password_salt
    new_user['password'] = new_user_password_encrypted
    new_user['study'] = True
    new_user['role'] = DEFAULT_USER_ROLE
    new_user['add_time'] = current_time_in_seconds
    new_user['up_time'] = current_time_in_seconds
    new_user['type'] = DEFAULT_REGISTER_TYPE
    new_user['--v'] = 0

    model_user.insert_one(new_user)
    print('INFO: create user successfully, id=%s' % str(new_user_id))

def prepare_arg_parser():
    arg_parser = argparse.ArgumentParser()
    arg_parser.add_argument('--mongo_host', dest='mongo_host', required=True)
    arg_parser.add_argument('--mongo_port', dest='mongo_port', required=True)
    arg_parser.add_argument('--username', dest='username', required=False)
    arg_parser.add_argument('--useremail', dest='useremail', required=True)
    return arg_parser


if __name__ == '__main__':
    arg_parser = prepare_arg_parser()
    args = arg_parser.parse_args()

    param_username = args.username
    param_useremail = args.useremail

    email_regexp_matcher = re.compile(REGEXP_EMAIL)
    if not email_regexp_matcher.match(param_useremail):
        print('ERROR: invalid email format, skip')
        exit(1)

    if param_username is None or len(param_username) == 0:
        param_username = param_useremail.split('@')[0]

    create_user(args.mongo_host, int(args.mongo_port), param_username, param_useremail)