diff --git a/app.py b/app.py index 81a5352..f18df2c 100644 --- a/app.py +++ b/app.py @@ -1,6 +1,8 @@ from flask import Flask, Response, request, session, abort, redirect from flask.logging import default_handler import flask +import flask.json +from werkzeug.exceptions import Unauthorized, Forbidden from werkzeug.middleware.proxy_fix import ProxyFix import uw_saml2 from urllib.parse import urljoin @@ -36,6 +38,10 @@ def configure_logging(): ) +def wants_json(request): + return "application/json" in request.accept_mimetypes.values() + + @app.route('/status') # if we add any more options then refactor all this. @app.route('/status/2fa') @app.route('/status/group/') @@ -56,16 +62,24 @@ def status(group=None): wants_2fa = '2fa' in request.path.split('/') has_2fa = session.get('has_2fa', False) if not userid or (wants_2fa and not has_2fa): - abort(401) + raise Unauthorized if group and group not in groups: message = f"{userid} not a member of {group} or SP can't receive it" app.logger.error(message) - abort(403) + raise Forbidden str_2fa = str(has_2fa).lower() headers = {'X-Saml-User': userid, 'X-Saml-Groups': ':'.join(groups), 'X-Saml-2fa': str_2fa} - txt = f'Logged in as: {userid}\nGroups: {str(groups)}\n2FA: {str_2fa}' + if wants_json(request): + txt = flask.json.dumps({ + "user": userid, + "groups": groups, + "two_factor": has_2fa + }) + headers['Content-Type'] = 'application/json' + else: + txt = f'Logged in as: {userid}\nGroups: {str(groups)}\n2FA: {str_2fa}' return Response(txt, status=200, headers=headers) @@ -139,3 +153,12 @@ def healthz():

Status

Logout

''' + + +@app.errorhandler(Unauthorized) +@app.errorhandler(Forbidden) +def error_handler(e): + if wants_json(request): + return flask.jsonify(error=str(e)), e.code + else: + return e