<reponame>ubwbu/cnlglsearch<gh_stars>0
import logging
import random
import time
import traceback
from datetime import datetime
from functools import wraps
from io import BytesIO

from captcha.image import ImageCaptcha
from cnlglsearch import db, mail
from cnlglsearch.models import User, UserInfo
from cnlglsearch.utils.access_manager import (AccountTokenRedis,
                                              ResetTokenRedis, SessionRedis)
from elasticsearch.client import logger
from flask import Blueprint, make_response, redirect, request
from flask.globals import g
from flask_mail import Message
from sqlalchemy.sql.sqltypes import String
from werkzeug.security import check_password_hash, generate_password_hash
from werkzeug.wrappers import Response

VERSION_AUTH = 'v0.0.1'

CODE_CHAR_RANGE = 'ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789'

bp = Blueprint('auth', __name__, url_prefix='/api/'+VERSION_AUTH)

logger = logging.getLogger(__name__)


@bp.before_app_request
def load_logged_in_user():
    g.sid = request.cookies.get('sid')
    g.sid = g.sid if g.sid else ""
    g.sid_older = g.sid
    g.uid = None
    g.username = None
    g.is_guest = True
    # 游客检索结果最大翻页限制
    g.max_page_limit = 10
    if not SessionRedis().check_lifetime():
        SessionRedis.create()
    if not g.sid:
        logger.error("create session failed!")
        return {"status": "error", "msg": "create session failed!"}, 500
    session = SessionRedis()
    g.uid = session.get("uid")
    g.username = session.get("username")
    # 若session 中存在 uid 和 username,
    # 但用户 access tokens 不存在该 sid，则根据 sid 删除该 session
    if g.uid:
        if g.sid in AccountTokenRedis(g.uid).tokens(g.sid):
            g.is_guest = False
            g.max_page_limit = 100
        else:   # 登录状态失效
            SessionRedis.delete_session()
            # return {"status":"redirect","msg":"redirect('/login')"}, 403
            return redirect('/login'), 403


def login_check(view):
    @wraps(view)
    def wrapped_view(**kwargs):
        if g.uid and g.sid and not g.is_guest:
            return view(**kwargs)
        resp = make_response({"staus": "failed", "msg": "Unauthorized!"}, 403)
        resp.delete_cookie('sid', '/')
        return resp
    return wrapped_view


# 使用 g.sid 判断是否存在 sid， 需配合 before_app_request
@bp.after_app_request
def set_cookies_session_id(response: Response):
    if g.sid != g.sid_older:
        response.set_cookie("sid", g.sid)
    return response


@bp.route('/login', methods=['POST'])
def login():
    """
    Email
    登录
    """
    result_dict = {"msg": None, "status": None}

    # 检查登录状态，重复登录则拒绝
    if g.sid and not g.is_guest:
        return {"status": "error",
                "msg": f'Hi {g.username}. You have already logged in!'}

    try:
        user_login = request.get_json()
    except Exception:
        logger.warning(traceback.format_exc())
        return {"stauts": "warning", "msg": "no data"}

    email = user_login.get('email')
    password = user_login.get('password')
    if not email or not password:
        return {"stauts": "warning", "msg": "no data"}
    q = {'email': email}
    # 数据库中查找用户
    user = User.query.filter_by(**q).first()
    # 密码校验成功
    if user is not None and check_password_hash('$'.join(
            [user.hash_method, user.salt, user.password]), password):
        username = user.username
        uid = user.uid
        user.last_login_ts = int(datetime.now().timestamp())
        db.session.commit()
        SessionRedis.new_session(uid=uid, username=username)
        sid = g.sid
        if sid:
            AccountTokenRedis(uid).add(
                sid, int(datetime.now().timestamp()))
            if sid not in AccountTokenRedis(uid).tokens(match=sid):
                return {'status': 'error', 'msg': 'failed'}, 500

        result_dict['msg'] = 'success'
        result_dict['status'] = 'ok'
        return result_dict
    else:
        result_dict['status'] = 'error'
        result_dict['msg'] = "User not exits or password error!"
    return result_dict


@bp.route('/logout', methods=['GET'])
def logout():
    """
    登出
    """
    if g.uid:
        actk = AccountTokenRedis(g.uid)
        if actk.ttl > 0:
            actk.remove(g.sid)
    SessionRedis.delete_session()
    return redirect('/index')


@bp.route('/signup', methods=['POST'])
def signup():
    """
    创建用户
    """

    usr_json = request.get_json()

    username = usr_json.get('username')
    email = usr_json.get('email')
    cap_code = usr_json.get("captcha_code")

    if not verify_captcha_code(cap_code):
        return {"status": "failed", "type": "captcha",
                "msg": "Verification code does not match"}
    if User.query.filter_by(username=username).first() is not None:
        return {"status": "failed", "type": "username", "msg": "username exists"}
    if User.query.filter_by(email=email).first() is not None:
        return {"status": "failed", "type": "email", "msg": "email exists"}

    password = <PASSWORD>('password')

    if not (username and email and password):
        return {"status": "error", "msg": "no data"}

    hash_method, salt, password = generate_password_hash(
        password=password).split('$')
    created_ts = int(datetime.now().timestamp())

    # uid
    count_rows = User.query.count()
    count_rows = count_rows - 1 if count_rows > 0 else 0
    u = User.query.offset(count_rows).first()
    if u is not None:
        uid = str(int(u.uid)+1)
    else:
        uid = 1

    new_user = User(
        uid=uid, username=username, email=email, hash_method=hash_method, salt=salt,
        password=password, created_ts=created_ts, updated_ts=None, last_login_ts=None
    )
    r = new_user.create()
    if isinstance(r, String):
        logger.error(r)
        return {'status': "error", "msg": "please try again"}, 500

    # 国家或地区id
    country_region_id = "39"
    try:
        country_region_id = int(usr_json.get("country_region_id"))
        if country_region_id > 0 and country_region_id < 190:
            country_region_id = str(country_region_id)
    except:
        ...
    userinfo = UserInfo(
        uid=uid,
        country_region_id=country_region_id,
        state=usr_json.get("state"),
        city=usr_json.get("city"),
        organization=usr_json.get("organization"),
        address=usr_json.get("address"),
        phone=usr_json.get("phone"),
        fist_name=usr_json.get("first_name"),
        last_name=usr_json.get("last_name"),
    )
    r = userinfo.create()
    if isinstance(r, String):
        logger.error(r)
    return {'status': 'ok'}


@bp.route('/reset_password', methods=['POST'])
def reset_passwd():
    req = request.get_json()
    reset_token = req.get("tk")
    resettk = ResetTokenRedis(reset_token)
    passwd = req.get("password") if req else None
    if not passwd:
        return {"status": "failed", "msg": "empty password"}, 200
    if resettk.check_lifetime():
        email = resettk["email"]
        if type(email) != str or email == "":
            return {"status": "error", "msg": "server error"}
        user = User.query.filter_by(email=email).first()
        if user is None:
            return {"status": "failed", "msg": "user not exits"}
        uid = user.uid
        user.hash_method, user.salt, user.password = generate_password_hash(
            password=<PASSWORD>).split('$', 2)
        user.updated_ts = int(datetime.now().timestamp())
        db.session.commit()
        ResetTokenRedis.delete_token(reset_token)
        if AccountTokenRedis(uid).ttl() > 0:
            AccountTokenRedis.delete_acc_token(uid)
        SessionRedis.delete_session()
        return {"status": "ok", "msg": "ok"}, 200
    else:
        return {"status": "failed", "msg": "reset token expired"}, 200


@bp.route("/captcha")
def captcha():
    session = SessionRedis()
    cap = "".join(random.sample(CODE_CHAR_RANGE, 6))
    session["captcha_code"] = cap.lower()
    session["captcha_exp_time"] = str(int(time.time())+60)
    im = ImageCaptcha().generate_image(cap)
    b = BytesIO()
    im.save(b, "webp")
    resp = make_response(b.getvalue())
    resp.headers["Content-Type"] = "image/webp"
    return resp


# {"s_type":"email", "email":"xx@xxx"}
@bp.route("/send_verify_code", methods=["POST"])
def send_verify_code():
    session = SessionRedis()
    v_code = "".join(random.sample(CODE_CHAR_RANGE, 8)).upper()

    req_data = request.get_json()
    s_type = req_data.get("s_type")
    if s_type == "email":
        mail_to = req_data.get("email")
        # 简单检测是否是邮箱
        if type(mail_to) != str or "@" not in mail_to:
            return {"status": "failed", "type": "email", "msg": "Check your email address."}
        # 检测用户是否注册
        if User.query.filter_by(email=mail_to).first() is None:
            return {"status": "failed", "type": "email", "msg": "User not exists"}

        msg = Message("Reset your account password", recipients=[mail_to], body="Please note that you are resetting your password! " +
                      "Your verification code is:  " + v_code)
        try:
            mail.send(msg)
            session["mail_verify"] = mail_to
            session["mail_verify_code"] = v_code
            session["mail_verify_code_exp_time"] = str(int(time.time())+60)
        except:
            logger.error(traceback.format_exc())
            return {"status": "error", "msg": "server error"}, 500
        return {"status": "ok", "msg": "ok"}
    return {"status": "error", "msg": "only support reset password by email"}


# {"s_type":"email", "email":"xx@xxx", "code":"xxx"}
@bp.route("/verify_code", methods=["POST"])
def verify_code():
    req = request.get_json()
    if req.get("s_type") == "email":
        code = req.get("code")
    email = req.get("email")
    flag, mail_e = verify_mail_code(code)
    if flag:
        if mail_e == email:
            tk = ResetTokenRedis.create(email)
            if tk:
                return {"status": "ok", "msg": tk}
    return {"status": "failed", "msg": "Code does not match!"}


def verify_captcha_code(code):
    session = SessionRedis()
    try:
        code = code.lower()
        exp_time = int(session["captcha_exp_time"])
        if exp_time > int(time.time()) and code and session["captcha_code"]:
            if session["captcha_code"] == code:
                return True
        else:
            session.delete("captcha_exp_time")
            session.delete("captcha_code")
        return False
    except:
        logger.error(traceback.format_exc())
        return False
    finally:
        ...


def verify_mail_code(code):
    """
    return (True, email) if successful else (False, None)
    """
    session = SessionRedis()
    try:
        exp_time = session["mail_verify_code_exp_time"]
        exp_time = int(exp_time) if exp_time else -100
        if exp_time > int(time.time()) and code and session["mail_verify_code"]:
            if session["mail_verify_code"] == code:
                email = session["mail_verify"]
                session.delete("mail_verify")
                session.delete("mail_verify_code_exp_time")
                session.delete("mail_verify_code")
                return True, email
        elif exp_time <= int(time.time()):
            session.delete("mail_verify")
            session.delete("mail_verify_code_exp_time")
            session.delete("mail_verify_code")
        return False, None
    except:
        logger.error(traceback.format_exc())
        return False, None

import logging
import random
import time
import traceback
from datetime import datetime
from functools import wraps
from io import BytesIO

from captcha.image import ImageCaptcha
from cnlglsearch import db, mail
from cnlglsearch.models import User, UserInfo
from cnlglsearch.utils.access_manager import (AccountTokenRedis,
                                              ResetTokenRedis, SessionRedis)
from elasticsearch.client import logger
from flask import Blueprint, make_response, redirect, request
from flask.globals import g
from flask_mail import Message
from sqlalchemy.sql.sqltypes import String
from werkzeug.security import check_password_hash, generate_password_hash
from werkzeug.wrappers import Response

VERSION_AUTH = 'v0.0.1'

CODE_CHAR_RANGE = 'ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789'

bp = Blueprint('auth', __name__, url_prefix='/api/'+VERSION_AUTH)

logger = logging.getLogger(__name__)


@bp.before_app_request
def load_logged_in_user():
    g.sid = request.cookies.get('sid')
    g.sid = g.sid if g.sid else ""
    g.sid_older = g.sid
    g.uid = None
    g.username = None
    g.is_guest = True
    # 游客检索结果最大翻页限制
    g.max_page_limit = 10
    if not SessionRedis().check_lifetime():
        SessionRedis.create()
    if not g.sid:
        logger.error("create session failed!")
        return {"status": "error", "msg": "create session failed!"}, 500
    session = SessionRedis()
    g.uid = session.get("uid")
    g.username = session.get("username")
    # 若session 中存在 uid 和 username,
    # 但用户 access tokens 不存在该 sid，则根据 sid 删除该 session
    if g.uid:
        if g.sid in AccountTokenRedis(g.uid).tokens(g.sid):
            g.is_guest = False
            g.max_page_limit = 100
        else:   # 登录状态失效
            SessionRedis.delete_session()
            # return {"status":"redirect","msg":"redirect('/login')"}, 403
            return redirect('/login'), 403


def login_check(view):
    @wraps(view)
    def wrapped_view(**kwargs):
        if g.uid and g.sid and not g.is_guest:
            return view(**kwargs)
        resp = make_response({"staus": "failed", "msg": "Unauthorized!"}, 403)
        resp.delete_cookie('sid', '/')
        return resp
    return wrapped_view


# 使用 g.sid 判断是否存在 sid， 需配合 before_app_request
@bp.after_app_request
def set_cookies_session_id(response: Response):
    if g.sid != g.sid_older:
        response.set_cookie("sid", g.sid)
    return response


@bp.route('/login', methods=['POST'])
def login():
    """
    Email
    登录
    """
    result_dict = {"msg": None, "status": None}

    # 检查登录状态，重复登录则拒绝
    if g.sid and not g.is_guest:
        return {"status": "error",
                "msg": f'Hi {g.username}. You have already logged in!'}

    try:
        user_login = request.get_json()
    except Exception:
        logger.warning(traceback.format_exc())
        return {"stauts": "warning", "msg": "no data"}

    email = user_login.get('email')
    password = user_login.get('password')
    if not email or not password:
        return {"stauts": "warning", "msg": "no data"}
    q = {'email': email}
    # 数据库中查找用户
    user = User.query.filter_by(**q).first()
    # 密码校验成功
    if user is not None and check_password_hash('$'.join(
            [user.hash_method, user.salt, user.password]), password):
        username = user.username
        uid = user.uid
        user.last_login_ts = int(datetime.now().timestamp())
        db.session.commit()
        SessionRedis.new_session(uid=uid, username=username)
        sid = g.sid
        if sid:
            AccountTokenRedis(uid).add(
                sid, int(datetime.now().timestamp()))
            if sid not in AccountTokenRedis(uid).tokens(match=sid):
                return {'status': 'error', 'msg': 'failed'}, 500

        result_dict['msg'] = 'success'
        result_dict['status'] = 'ok'
        return result_dict
    else:
        result_dict['status'] = 'error'
        result_dict['msg'] = "User not exits or password error!"
    return result_dict


@bp.route('/logout', methods=['GET'])
def logout():
    """
    登出
    """
    if g.uid:
        actk = AccountTokenRedis(g.uid)
        if actk.ttl > 0:
            actk.remove(g.sid)
    SessionRedis.delete_session()
    return redirect('/index')


@bp.route('/signup', methods=['POST'])
def signup():
    """
    创建用户
    """

    usr_json = request.get_json()

    username = usr_json.get('username')
    email = usr_json.get('email')
    cap_code = usr_json.get("captcha_code")

    if not verify_captcha_code(cap_code):
        return {"status": "failed", "type": "captcha",
                "msg": "Verification code does not match"}
    if User.query.filter_by(username=username).first() is not None:
        return {"status": "failed", "type": "username", "msg": "username exists"}
    if User.query.filter_by(email=email).first() is not None:
        return {"status": "failed", "type": "email", "msg": "email exists"}

    password = <PASSWORD>('password')

    if not (username and email and password):
        return {"status": "error", "msg": "no data"}

    hash_method, salt, password = generate_password_hash(
        password=password).split('$')
    created_ts = int(datetime.now().timestamp())

    # uid
    count_rows = User.query.count()
    count_rows = count_rows - 1 if count_rows > 0 else 0
    u = User.query.offset(count_rows).first()
    if u is not None:
        uid = str(int(u.uid)+1)
    else:
        uid = 1

    new_user = User(
        uid=uid, username=username, email=email, hash_method=hash_method, salt=salt,
        password=password, created_ts=created_ts, updated_ts=None, last_login_ts=None
    )
    r = new_user.create()
    if isinstance(r, String):
        logger.error(r)
        return {'status': "error", "msg": "please try again"}, 500

    # 国家或地区id
    country_region_id = "39"
    try:
        country_region_id = int(usr_json.get("country_region_id"))
        if country_region_id > 0 and country_region_id < 190:
            country_region_id = str(country_region_id)
    except:
        ...
    userinfo = UserInfo(
        uid=uid,
        country_region_id=country_region_id,
        state=usr_json.get("state"),
        city=usr_json.get("city"),
        organization=usr_json.get("organization"),
        address=usr_json.get("address"),
        phone=usr_json.get("phone"),
        fist_name=usr_json.get("first_name"),
        last_name=usr_json.get("last_name"),
    )
    r = userinfo.create()
    if isinstance(r, String):
        logger.error(r)
    return {'status': 'ok'}


@bp.route('/reset_password', methods=['POST'])
def reset_passwd():
    req = request.get_json()
    reset_token = req.get("tk")
    resettk = ResetTokenRedis(reset_token)
    passwd = req.get("password") if req else None
    if not passwd:
        return {"status": "failed", "msg": "empty password"}, 200
    if resettk.check_lifetime():
        email = resettk["email"]
        if type(email) != str or email == "":
            return {"status": "error", "msg": "server error"}
        user = User.query.filter_by(email=email).first()
        if user is None:
            return {"status": "failed", "msg": "user not exits"}
        uid = user.uid
        user.hash_method, user.salt, user.password = generate_password_hash(
            password=<PASSWORD>).split('$', 2)
        user.updated_ts = int(datetime.now().timestamp())
        db.session.commit()
        ResetTokenRedis.delete_token(reset_token)
        if AccountTokenRedis(uid).ttl() > 0:
            AccountTokenRedis.delete_acc_token(uid)
        SessionRedis.delete_session()
        return {"status": "ok", "msg": "ok"}, 200
    else:
        return {"status": "failed", "msg": "reset token expired"}, 200


@bp.route("/captcha")
def captcha():
    session = SessionRedis()
    cap = "".join(random.sample(CODE_CHAR_RANGE, 6))
    session["captcha_code"] = cap.lower()
    session["captcha_exp_time"] = str(int(time.time())+60)
    im = ImageCaptcha().generate_image(cap)
    b = BytesIO()
    im.save(b, "webp")
    resp = make_response(b.getvalue())
    resp.headers["Content-Type"] = "image/webp"
    return resp


# {"s_type":"email", "email":"xx@xxx"}
@bp.route("/send_verify_code", methods=["POST"])
def send_verify_code():
    session = SessionRedis()
    v_code = "".join(random.sample(CODE_CHAR_RANGE, 8)).upper()

    req_data = request.get_json()
    s_type = req_data.get("s_type")
    if s_type == "email":
        mail_to = req_data.get("email")
        # 简单检测是否是邮箱
        if type(mail_to) != str or "@" not in mail_to:
            return {"status": "failed", "type": "email", "msg": "Check your email address."}
        # 检测用户是否注册
        if User.query.filter_by(email=mail_to).first() is None:
            return {"status": "failed", "type": "email", "msg": "User not exists"}

        msg = Message("Reset your account password", recipients=[mail_to], body="Please note that you are resetting your password! " +
                      "Your verification code is:  " + v_code)
        try:
            mail.send(msg)
            session["mail_verify"] = mail_to
            session["mail_verify_code"] = v_code
            session["mail_verify_code_exp_time"] = str(int(time.time())+60)
        except:
            logger.error(traceback.format_exc())
            return {"status": "error", "msg": "server error"}, 500
        return {"status": "ok", "msg": "ok"}
    return {"status": "error", "msg": "only support reset password by email"}


# {"s_type":"email", "email":"xx@xxx", "code":"xxx"}
@bp.route("/verify_code", methods=["POST"])
def verify_code():
    req = request.get_json()
    if req.get("s_type") == "email":
        code = req.get("code")
    email = req.get("email")
    flag, mail_e = verify_mail_code(code)
    if flag:
        if mail_e == email:
            tk = ResetTokenRedis.create(email)
            if tk:
                return {"status": "ok", "msg": tk}
    return {"status": "failed", "msg": "Code does not match!"}


def verify_captcha_code(code):
    session = SessionRedis()
    try:
        code = code.lower()
        exp_time = int(session["captcha_exp_time"])
        if exp_time > int(time.time()) and code and session["captcha_code"]:
            if session["captcha_code"] == code:
                return True
        else:
            session.delete("captcha_exp_time")
            session.delete("captcha_code")
        return False
    except:
        logger.error(traceback.format_exc())
        return False
    finally:
        ...


def verify_mail_code(code):
    """
    return (True, email) if successful else (False, None)
    """
    session = SessionRedis()
    try:
        exp_time = session["mail_verify_code_exp_time"]
        exp_time = int(exp_time) if exp_time else -100
        if exp_time > int(time.time()) and code and session["mail_verify_code"]:
            if session["mail_verify_code"] == code:
                email = session["mail_verify"]
                session.delete("mail_verify")
                session.delete("mail_verify_code_exp_time")
                session.delete("mail_verify_code")
                return True, email
        elif exp_time <= int(time.time()):
            session.delete("mail_verify")
            session.delete("mail_verify_code_exp_time")
            session.delete("mail_verify_code")
        return False, None
    except:
        logger.error(traceback.format_exc())
        return False, None
# 游客检索结果最大翻页限制 # 若session 中存在 uid 和 username, # 但用户 access tokens 不存在该 sid，则根据 sid 删除该 session # 登录状态失效 # return {"status":"redirect","msg":"redirect('/login')"}, 403 # 使用 g.sid 判断是否存在 sid， 需配合 before_app_request Email     登录 # 检查登录状态，重复登录则拒绝 # 数据库中查找用户 # 密码校验成功 登出 创建用户 # uid # 国家或地区id # {"s_type":"email", "email":"xx@xxx"} # 简单检测是否是邮箱 # 检测用户是否注册 # {"s_type":"email", "email":"xx@xxx", "code":"xxx"} return (True, email) if successful else (False, None)