107 lines
4.9 KiB
Python
107 lines
4.9 KiB
Python
from tools import *
|
|
from Utils.Config import *
|
|
from Utils.Lib import *
|
|
from Utils.Models import *
|
|
from AppBridge import *
|
|
|
|
GERP: flask.Blueprint = Others["GERP"]
|
|
|
|
|
|
|
|
@GERP.route('/setup-2fa', methods=['GET', 'POST'])
|
|
def Setup2fa():
|
|
User = GetELTUser()
|
|
if not User:
|
|
flask.flash('请先登录', 'error')
|
|
return flask.redirect(flask.url_for('login'))
|
|
Existing2fa = GERPUser2FA.query.filter_by(UserUUID=User.UserUUID).first()
|
|
if flask.request.method == 'POST':
|
|
if 'generate' in flask.request.form:
|
|
Secret = pyotp.random_base32()
|
|
flask.session['2fa_secret'] = Secret
|
|
flask.session['2fa_user_uuid'] = User.UserUUID.hex()
|
|
Totp = pyotp.TOTP(Secret)
|
|
Username = User.UserName + " / 银河锋锐软件开发工作室"
|
|
IssuerName = "GVSDS"
|
|
QrUrl = Totp.provisioning_uri(Username, issuer_name=IssuerName)
|
|
QrImg = qrcode.make(QrUrl)
|
|
ImgBuffer = BytesIO()
|
|
QrImg.save(ImgBuffer, format='PNG')
|
|
ImgBuffer.seek(0)
|
|
QrCode = base64.b64encode(ImgBuffer.getvalue()).decode('utf-8')
|
|
return flask.render_template('/setup2fa.html',
|
|
user=User,
|
|
qr_code=QrCode,
|
|
secret=Secret,
|
|
step=2)
|
|
elif 'verify' in flask.request.form:
|
|
if '2fa_secret' not in flask.session or '2fa_user_uuid' not in flask.session:
|
|
flask.flash('会话已过期，请重新开始设置', 'error')
|
|
return flask.redirect(flask.url_for('GERP.Setup2fa'))
|
|
if flask.session['2fa_user_uuid'] != User.UserUUID.hex():
|
|
flask.flash('验证失败，用户不匹配', 'error')
|
|
return flask.redirect(flask.url_for('GERP.Setup2fa'))
|
|
Token = flask.request.form.get('token', '').strip()
|
|
Totp = pyotp.TOTP(flask.session['2fa_secret'])
|
|
if not Token.isdigit() or len(Token) != 6:
|
|
flask.flash('验证码必须是6位数字', 'error')
|
|
return flask.render_template('/setup2fa.html', step=2, secret=flask.session['2fa_secret'])
|
|
if Totp.verify(Token, valid_window=1):
|
|
if Existing2fa:
|
|
Existing2fa.SecretKey = flask.session['2fa_secret']
|
|
Existing2fa.IsEnabled = 1
|
|
else:
|
|
New2fa = GERPUser2FA(
|
|
UserUUID=User.UserUUID,
|
|
SecretKey=flask.session['2fa_secret'],
|
|
IsEnabled=1
|
|
)
|
|
db.session.add(New2fa)
|
|
db.session.commit()
|
|
flask.session.pop('2fa_secret', None)
|
|
flask.session.pop('2fa_user_uuid', None)
|
|
flask.flash('2FA认证已成功设置', 'success')
|
|
return flask.redirect("/")
|
|
else:
|
|
flask.flash('验证码无效，请重试', 'error')
|
|
return flask.render_template('setup2fa.html', step=2, secret=flask.session['2fa_secret'])
|
|
elif 'disable' in flask.request.form:
|
|
if Existing2fa:
|
|
Existing2fa.IsEnabled = 0
|
|
db.session.commit()
|
|
flask.flash('2FA认证已禁用', 'success')
|
|
else:
|
|
flask.flash('您尚未设置2FA认证', 'error')
|
|
return flask.redirect(flask.url_for('GERP.Setup2fa'))
|
|
Secret = flask.session.get('2fa_secret')
|
|
Step = 2 if Secret else 1
|
|
return flask.render_template('setup2fa.html',
|
|
user=User,
|
|
existing2fa=Existing2fa,
|
|
step=Step,
|
|
secret=Secret)
|
|
|
|
@GERP.route('/verify-2fa', methods=['GET', 'POST'])
|
|
def Verify2fa():
|
|
User = GetELTUser()
|
|
if not User:
|
|
flask.flash('请先登录', 'error')
|
|
return flask.redirect(flask.url_for('login'))
|
|
User2fa = GERPUser2FA.query.filter_by(UserUUID=User.UserUUID, IsEnabled=1).first()
|
|
if not User2fa:
|
|
flask.flash('您尚未设置2FA认证，请先完成设置', 'error')
|
|
return flask.redirect(flask.url_for('GERP.Setup2fa'))
|
|
if flask.request.method == 'POST':
|
|
Token = flask.request.form.get('token', '').strip()
|
|
if not Token.isdigit() or len(Token) != 6:
|
|
flask.flash('验证码必须是6位数字', 'error')
|
|
return flask.render_template('verify2fa.html', user=User)
|
|
Totp = pyotp.TOTP(User2fa.SecretKey)
|
|
if Totp.verify(Token, valid_window=1):
|
|
User2fa.LastUsedAt = datetime.datetime.now(Timezone_UTFA8)
|
|
db.session.commit()
|
|
flask.flash('2FA验证成功', 'success')
|
|
else:
|
|
flask.flash('2FA验证失败，验证码无效', 'error')
|
|
return flask.render_template('verify2fa.html', user=User)
|