import json from argparse import ArgumentParser from operator import attrgetter from os.path import join from flask import ( Flask, render_template, request, redirect, url_for, send_from_directory, flash, make_response, send_file, session, ) from flask_compress import Compress from flask_login import ( LoginManager, login_required, logout_user, login_user, current_user, ) from config import config from user import ( UserFactory, AUTH_TOKEN_COOKIE_KEY, CreateUserError, ) log = config.logger(__file__) log.info("################ Web Server initializing #####################") app = Flask(__name__) app.config.from_mapping(config.get("webserver.flask")) app.config.update( SECRET_KEY=config["secure.http.session_secret.webserver"], COMPRESS_MIMETYPES=[ "text/html", "text/css", "text/xml", "application/json", "application/javascript", "text/javascript", ], COMPRESS_LEVEL=9, ) login_manager = LoginManager(app) login_manager.login_view = "login" Compress(app) def _secure_url_for(endpoint, external=True, **values): if not config.get("webserver.redirect_to_https", False): return url_for(endpoint, _external=external, **values) return url_for(endpoint, _external=external, _scheme="https", **values) @login_manager.user_loader def load_user(id): return UserFactory.get_class().get(id) @login_manager.unauthorized_handler def handle_needs_login(): session.pop("_flashes", None) flash("You have to be logged in to access this page.") return redirect( _secure_url_for(".login", next=request.endpoint, **request.view_args) ) @app.route("/create_user", methods=["GET", "POST"]) def create_user(): data = request.args or request.form name = data["name"].strip() try: user = UserFactory.get_class().create_by_name(name) except CreateUserError as e: session.pop("_flashes", None) message = e.args[0] if "value combination" in message.lower(): message = f"Failed creating user {name}" flash(message) return redirect(_secure_url_for(".login")) return _complete_user_login(user) @app.route("/login/") def login_by_id(user_id): if current_user.get_id() == user_id and current_user.is_authenticated: return redirect(_secure_url_for(".login")) try: user = load_user(user_id) except Exception as e: # Some callback issue, try again. # For example, the user tried to reload the callback page and the provider token was already redeemed. session.pop("_flashes", None) if e.args: flash(e.args[0]) else: flash(repr(e)) return redirect(_secure_url_for(".login")) args = dict(request.args) endpoint = args.pop("next", ".index") return _complete_user_login(user, endpoint, **request.args) def _complete_user_login(user, endpoint=".index", **kwargs): login_user(user, True) response = redirect(_secure_url_for(endpoint, **kwargs)) set_response_cookie(response, user) return response def set_response_cookie(response, user, copy_request=None): if copy_request and AUTH_TOKEN_COOKIE_KEY in copy_request.cookies: token = request.cookies[AUTH_TOKEN_COOKIE_KEY] else: token = user.token response.set_cookie( AUTH_TOKEN_COOKIE_KEY, token, **config.get("webserver.auth.cookies"), ) @app.route("/logout") def logout(): logout_user() return redirect(_secure_url_for(".login")) @app.route("/static/network/") @login_required def send_netron(): return send_from_directory("static", join("network", "index.html")) @app.route("/static/") @login_required def send_static(path): return send_from_directory("static", path) @app.route("/webapp_conf.js") def webapp_conf(): webapp_conf = _get_webapp_conf() response = make_response(f"SM_CONFIG={json.dumps(webapp_conf)}") response.content_type = "text/javascript" return response def _get_webapp_conf(): webapp_conf = config.get("webapp") webapp_build = None webapp_path = "pages/webapp" try: with open(f"{webapp_path}/assets/build.json") as data_file: webapp_build = json.load(data_file) webapp_build["branch"] = webapp_build["branch"].replace("@", "") except Exception: webapp_build = {} backend_path = "." try: with open(f"{backend_path}/static/build.json") as data_file: backend_build = json.load(data_file) except Exception: backend_build = {} webapp_conf.put("backend_build", backend_build) webapp_conf.put("webapp_build", webapp_build) webapp_conf.put("env", config.env) return webapp_conf @app.route("/login") def login(): users = sorted((user.user_data for user in UserFactory.get_class().get_all()), key=attrgetter("name")) response = make_response(render_template("login.html", users=users)) # make sure to clear out session token cookie response.set_cookie( AUTH_TOKEN_COOKIE_KEY, "", expires=0, **config.get("webserver.auth.cookies") ) return response def _serve_webapp(path=None): if not path: response = make_response(send_file("pages/webapp/index.html")) else: try: response = make_response(send_from_directory("pages/webapp", path)) except Exception: response = make_response(send_file("pages/webapp/index.html")) set_response_cookie(response, current_user, request) response.headers["X-Trains-Environment"] = "oss" return response @app.route("/favicon.ico") def favicon(): return send_from_directory("static", "favicon.ico") @app.route("/") def index(): if not current_user.is_authenticated: return redirect(_secure_url_for(".login")) return _serve_webapp() @app.route("/") @login_required def webapp(path=None): return _serve_webapp(path) def parse_args(): parser = ArgumentParser(description=__doc__) parser.add_argument( "--port", "-p", type=int, default=config.get("webserver.listen.port") ) parser.add_argument("--ip", "-i", type=str, default=config.get("webserver.listen.ip")) parser.add_argument( "--debug", action="store_true", default=config.get("webserver.debug") ) return parser.parse_args() def main(): args = parse_args() app.run(debug=args.debug, host=args.ip, port=args.port, threaded=True) if __name__ == "__main__": main()