from PDAPI import *

TOTP_CONFIG = {
    "digits": 6,
    "interval": 30,
}


__metadata__ = {
    "methods": ["GET"],
    "blueprint": "GERP",
    "permission": 1.0,
    "attributes": [],
    "limit": "1 per second",
}

def HandleRequest(_CleanedData):
    current_user = UGM.GetELTUser()
    if not current_user:
        return {"code": 401, "msg": "未登录"}, 401

    existing = GERPUser2FA.query.filter(GERPUser2FA.UserUUID == current_user.UserUUID).first()

    if existing and existing.IsEnabled == 1:
        return {"code": 1, "msg": "2FA已启用，请先禁用后再重新设置"}, 200

    Secret = pyotp.random_base32()
    Totp = pyotp.TOTP(Secret, digits=TOTP_CONFIG["digits"], interval=TOTP_CONFIG["interval"])
    QrUrl = Totp.provisioning_uri(
        f"{current_user.UserName} / {current_user.UserEmail}",
        issuer_name=f"GVSDS - {current_user.UserName}",
        image="https://www.gvsds.com/assets/hlogos/nobg.png"
    )
    QrImg = qrcode.make(QrUrl)
    ImgBuffer = BytesIO()
    QrImg.save(ImgBuffer, format='PNG')
    ImgBuffer.seek(0)

    if existing:
        existing.SecretKey = Secret
        existing.Verified = 0
        existing.IsEnabled = 0
    else:
        new_2fa = GERPUser2FA()
        new_2fa.User2FAUUID = NewBUUID()
        new_2fa.UserUUID = current_user.UserUUID
        new_2fa.SecretKey = Secret
        new_2fa.Verified = 0
        new_2fa.IsEnabled = 0
        db.session.add(new_2fa)

    try:
        db.session.commit()
    except:
        db.session.rollback()
        return {"code": 1, "msg": "SQL 提交失败"}, 200

    return flask.Response(ImgBuffer.getvalue(), mimetype='image/png')