from __future__ import annotations import os import hashlib import secrets from datetime import timedelta from flask import ( Blueprint, abort, current_app, flash, jsonify, redirect, render_template, request, send_file, session, url_for, ) from flask_login import current_user, login_required, login_user, logout_user from sqlalchemy.orm import joinedload from .config import BASE_DIR from .extensions import db from .models import AppSetting, PasswordResetToken, UploadRecord, User from .prediction import PredictionError, run_prediction from .security import new_captcha from .time_utils import format_datetime_for_timezone, utc_now bp = Blueprint("main", __name__) REFERENCE_PDF_NAME = "20260630标准文本——供水管道健康状态与剩余寿命评估技术导则.pdf" TEMPLATE_EXCEL_NAME = "管道预测数据模板.xlsx" REGISTRATION_SETTING_KEY = "allow_registration" RECORDS_PER_PAGE = 10 def render_auth_template(mode: str, status_code: int = 200, captcha: str = ""): return render_template("login.html", mode=mode, captcha=captcha), status_code def refresh_captcha() -> str: session["captcha"] = new_captcha() return session["captcha"] def render_auth_error(mode: str, message: str, status_code: int = 400): flash(message, "error") return render_auth_template(mode, status_code, refresh_captcha()) def captcha_is_valid() -> bool: captcha_input = request.form.get("captcha", "").strip().upper() return bool(captcha_input and captcha_input == session.get("captcha", "")) def password_reset_token_hash(token: str) -> str: return hashlib.sha256(token.encode("utf-8")).hexdigest() def password_reset_expiry() -> datetime: minutes = max(int(current_app.config["PASSWORD_RESET_TOKEN_MINUTES"]), 1) return utc_now() + timedelta(minutes=minutes) def format_app_datetime(value: datetime | None) -> str: return format_datetime_for_timezone(value, current_app.config["APP_TIMEZONE"]) def active_password_reset_token(token: str) -> PasswordResetToken | None: reset_token = PasswordResetToken.query.filter_by( token_hash=password_reset_token_hash(token) ).first() if ( reset_token is None or reset_token.used_at is not None or reset_token.expires_at <= utc_now() ): return None return reset_token def render_password_reset_unavailable(status_code: int = 400): flash("重置链接无效或已过期,请联系管理员重新生成。", "error") return render_template( "password_reset.html", reset_token=None, token="", token_available=False, ), status_code def require_admin() -> None: if not current_user.is_admin: abort(403) def registration_allowed() -> bool: return AppSetting.get_bool( REGISTRATION_SETTING_KEY, current_app.config["ALLOW_REGISTRATION"], ) def requested_page() -> int: try: return max(int(request.args.get("page", 1)), 1) except (TypeError, ValueError): return 1 def paginated_uploads(query, endpoint: str): page = requested_page() pagination = ( query.order_by(UploadRecord.upload_time.desc()) .paginate(page=page, per_page=RECORDS_PER_PAGE, error_out=False) ) if pagination.pages and page > pagination.pages: return pagination, redirect(url_for(endpoint, page=pagination.pages)) return pagination, None def password_reset_users(): return ( User.query.filter(User.is_admin.is_(False)) .order_by(User.username.asc(), User.id.asc()) .all() ) def prediction_result_payload(artifacts, record: UploadRecord) -> dict: image_url = url_for("static", filename=f"images/{artifacts.image_filename}") importance_url = ( url_for("static", filename=f"images/{artifacts.importance_filename}") if artifacts.importance_filename else None ) return { "original_filename": artifacts.original_filename, "generated_at": format_app_datetime(record.upload_time), "image_url": image_url, "importance_url": importance_url, "excel_url": url_for("main.download_file", record_id=record.id, file_type="prediction"), "result_url": url_for("main.result_page"), "sample_count": int(artifacts.sample_count), "summary_rows": artifacts.summary_rows[:6], "analysis_text": artifacts.analysis_text, } @bp.route("/") def index(): if current_user.is_authenticated: return redirect(url_for("main.home")) return redirect(url_for("main.login")) @bp.route("/login", methods=["GET", "POST"]) def login(): if current_user.is_authenticated: return redirect(url_for("main.home")) if request.method == "GET": return render_auth_template("login", captcha=refresh_captcha()) username = request.form.get("username", "").strip() password = request.form.get("password", "") if not captcha_is_valid(): return render_auth_error("login", "验证码错误") user = User.query.filter_by(username=username).first() if not user or not user.check_password(password): return render_auth_error("login", "用户名或密码错误") login_user(user, remember=bool(request.form.get("remember"))) return redirect(url_for("main.home")) @bp.route("/register", methods=["GET", "POST"]) def register(): if request.method == "GET": return render_auth_template("register", captcha=refresh_captcha()) username = request.form.get("username", "").strip() password = request.form.get("password", "") if not captcha_is_valid(): return render_auth_error("register", "验证码错误") if not registration_allowed(): return render_auth_error("register", "当前未开放自助注册,请联系管理员。", 403) if not username: return render_auth_error("register", "用户名不能为空") if len(password) < 6: return render_auth_error("register", "密码至少需要 6 位") if User.query.filter_by(username=username).first(): return render_auth_error("register", "用户名已存在") user = User(username=username, is_admin=False) user.set_password(password) db.session.add(user) db.session.commit() flash("注册成功,请登录", "info") return render_auth_template("login", captcha=refresh_captcha()) @bp.route("/password-reset/", methods=["GET", "POST"]) def password_reset(token: str): reset_token = active_password_reset_token(token) if reset_token is None: return render_password_reset_unavailable() if request.method == "GET": return render_template( "password_reset.html", reset_token=reset_token, token=token, token_available=True, ) password = request.form.get("password", "") password_confirm = request.form.get("password_confirm", "") if len(password) < 6: flash("密码至少需要 6 位", "error") return render_template( "password_reset.html", reset_token=reset_token, token=token, token_available=True, ), 400 if password != password_confirm: flash("两次输入的密码不一致", "error") return render_template( "password_reset.html", reset_token=reset_token, token=token, token_available=True, ), 400 reset_token.user.set_password(password) reset_token.used_at = utc_now() db.session.commit() flash("密码已重置,请使用新密码登录", "info") return render_auth_template("login", captcha=refresh_captcha()) @bp.route("/logout", methods=["POST"]) @login_required def logout(): logout_user() return redirect(url_for("main.login")) @bp.route("/home") @login_required def home(): return render_template("home.html") @bp.route("/history") @login_required def history_page(): pagination, page_redirect = paginated_uploads( UploadRecord.query.filter_by(user_id=current_user.id), "main.history_page", ) if page_redirect: return page_redirect return render_template( "history.html", pagination=pagination, records=pagination.items, ) @bp.route("/admin") @login_required def admin_dashboard(): require_admin() pagination, page_redirect = paginated_uploads( UploadRecord.query.options(joinedload(UploadRecord.user)), "main.admin_dashboard", ) if page_redirect: return page_redirect return render_template( "admin.html", pagination=pagination, records=pagination.items, password_reset_users=password_reset_users(), registration_allowed=registration_allowed(), ) @bp.route("/admin/registration", methods=["POST"]) @login_required def update_registration_setting(): require_admin() allow_registration = request.form.get("allow_registration") == "on" AppSetting.set_bool(REGISTRATION_SETTING_KEY, allow_registration) db.session.commit() message = "已开放用户自助注册" if allow_registration else "已关闭用户自助注册" if request.headers.get("X-Requested-With") == "XMLHttpRequest": return jsonify( { "message": message, "registration_allowed": allow_registration, "status_label": "已开放" if allow_registration else "已关闭", } ) flash(message, "info") return redirect(url_for("main.admin_dashboard")) @bp.route("/admin/users//password-reset-link", methods=["POST"]) @login_required def create_password_reset_link(user_id: int): require_admin() user = db.session.get(User, user_id) if user is None: return jsonify({"error": "用户不存在"}), 404 if user.is_admin: return jsonify({"error": "管理员账号不支持通过此入口重置密码"}), 403 now = utc_now() PasswordResetToken.query.filter_by(user_id=user.id, used_at=None).update( {"used_at": now} ) token = secrets.token_urlsafe(32) reset_token = PasswordResetToken( user_id=user.id, created_by_id=current_user.id, token_hash=password_reset_token_hash(token), expires_at=password_reset_expiry(), ) db.session.add(reset_token) db.session.commit() return jsonify( { "message": f"已生成 {user.username} 的密码重置链接", "reset_url": url_for("main.password_reset", token=token, _external=True), "expires_at": format_app_datetime(reset_token.expires_at), "user_id": user.id, "username": user.username, } ) @bp.route("/download//") @login_required def download_file(record_id: int, file_type: str): record = db.session.get(UploadRecord, record_id) if record is None: abort(404) if not (current_user.is_admin or current_user.id == record.user_id): abort(403) file_path = { "original": record.saved_path, "prediction": record.prediction_path, }.get(file_type) if file_path is None: abort(404) if not os.path.exists(file_path): abort(404) return send_file(file_path, as_attachment=True) @bp.route("/download_template") def download_template(): template_path = BASE_DIR / "example.xlsx" if not template_path.exists(): abort(404) return send_file(template_path, as_attachment=True, download_name=TEMPLATE_EXCEL_NAME) @bp.route("/reference_pdf") @login_required def reference_pdf(): pdf_path = BASE_DIR / REFERENCE_PDF_NAME if not pdf_path.exists(): abort(404) return send_file( pdf_path, as_attachment=False, download_name=REFERENCE_PDF_NAME, mimetype="application/pdf", ) @bp.route("/reference") @login_required def reference_page(): return render_template("reference.html") @bp.route("/result") @login_required def result_page(): result = session.get("last_result") return render_template("result.html", result=result) @bp.route("/predict", methods=["POST"]) @login_required def predict(): model = current_app.config.get("RSF_MODEL") if model is None: return jsonify({"error": "模型未成功加载,请检查模型文件。"}), 500 uploaded = request.files.get("file") if uploaded is None or uploaded.filename == "": return jsonify({"error": "未选择文件"}), 400 try: artifacts = run_prediction(uploaded, int(current_user.id), model) except PredictionError as exc: return jsonify({"error": exc.message}), exc.status_code record = UploadRecord( user_id=current_user.id, original_filename=artifacts.original_filename, saved_path=str(artifacts.saved_path), prediction_path=str(artifacts.excel_path), image_path=str(artifacts.image_path), ) db.session.add(record) db.session.commit() last_result = prediction_result_payload(artifacts, record) session["last_result"] = last_result return jsonify( { "message": "预测成功", "image_url": last_result["image_url"], "importance_url": last_result["importance_url"], "excel_url": last_result["excel_url"], "result_url": last_result["result_url"], "sample_count": last_result["sample_count"], "original_filename": artifacts.original_filename, } )