import os
import pyotp
import random
import os.path
import pathlib
from qrcode import QRCode, constants

# 获取随机密钥
# 你可以将该密钥存放在用户表中，嗯，记得加密
gtoken = pyotp.random_base32(64)
print('private key:',gtoken)

gtoken = 'TW2EPDI4NZE2XXL5Y3AXFQZUQVPCBMC3HRTGN57UI3AYAZLFBSGUH7S74JSGMONO'

# 然后根据密码创建二维码
def gen_scan_data(secret_key,username,issue):
    return pyotp.totp.TOTP(secret_key).provisioning_uri(username, issuer_name=issue)


data = gen_scan_data(gtoken,'yufei','简单教程')
print('scan data',data)

# 返回两个参数，成功与否和成功之后的二维码图片路径
def get_qrcode(data,username):
    BASE_DIR = os.path.dirname(os.path.abspath(__file__))
    dirpath = os.path.join(BASE_DIR, 'mfa', 'static', 'image')
    filepath = dirpath + os.sep + username + '.png'

    if os.path.isfile(filepath):
        return True

    if not pathlib.Path(dirpath).is_dir():
        os.makedirs(dirpath)

    qr = QRCode(
        version=1,
        error_correction=constants.ERROR_CORRECT_L,
        box_size=6,
        border=4, )
    try:
        qr.add_data(data)
        qr.make(fit=True)
        img = qr.make_image()
        img.save(filepath)  # 保存条形码图片
        return True,filepath
    except Exception as e:
        return False,None

# 检查用户输入的是否正确
# 返回 True / False
def google_authenticator_verify(secret_key, code):
    t = pyotp.TOTP(secret_key)
    return t.verify(code)

# 恢复代码
def gen_key(n=16):
    """
    获取密钥 n 密钥长度
    :return:
    """
    c_length = int(n)
    source = 'ABCDEFGHJKMNPQRSTWXYZ&*^abcdefhijkmnprstwxyz2345678@!$%~^&*_()'
    length = len(source) - 1
    result = ''
    for i in range(c_length):
        result += source[random.randint(0, length)]
    return result

def gen_recovery_code(cnt=10):
    return [ gen_key() for i in range(10)]


# 一般情况下使用客户端的 qrcode 来生成，这样可以减轻服务器的负担
print(get_qrcode(data,'yufei'))




# 验证
print('google_authenticator_verify:',google_authenticator_verify(gtoken,'340098'))



## 注意，恢复代码和 otp 是没有任何关系的，一般需要保存到数据库中，而且用一次就要删掉
## 创建恢复代码
## 你可以将该密钥存放在用户表中，嗯，记得加密
print("\n".join(gen_recovery_code()))