'''This version is intended to be run in development mode, using
copies of our CER and KEY files for Tempest. It needs to run on port
1942, since that's the only one that Doug has configured for us.

December 6, 2024
Scott D. Anderson
'''


import os
import logging

from flask import Flask, request, render_template, redirect, session, make_response, url_for

from onelogin.saml2.auth import OneLogin_Saml2_Auth
from onelogin.saml2.utils import OneLogin_Saml2_Utils


app = Flask(__name__)

app.logger.setLevel(logging.DEBUG)
handler = logging.FileHandler('./saml_ssl.log')
app.logger.addHandler(handler)

app.config["SECRET_KEY"] = "onelogindemopytoolkit"
app.config["SAML_PATH"] = os.path.join(os.path.dirname(os.path.abspath(__file__)), "saml")


def init_saml_auth(req):
    auth = OneLogin_Saml2_Auth(req, custom_base_path=app.config["SAML_PATH"])
    return auth


def prepare_flask_request(request):
    # If server is behind proxys or balancers use the HTTP_X_FORWARDED fields
    return {
        "https": "on" if request.scheme == "https" else "off",
        "http_host": request.host,
        "script_name": request.path,
        "get_data": request.args.copy(),
        # Uncomment if using ADFS as IdP, https://github.com/onelogin/python-saml/pull/144
        # 'lowercase_urlencoding': True,
        "post_data": request.form.copy(),
    }

@app.route("/", methods=["GET", "POST"])
def index():
    app.logger.info(f'SCOTTSAML access to / {request.method} and {request.args}')
    req = prepare_flask_request(request)
    auth = init_saml_auth(req)
    errors = []
    error_reason = None
    not_auth_warn = False
    success_slo = False
    attributes = False
    paint_logout = False

    if "sso" in request.args:
        ## SCOTT: this is where we transfer to the IDP (Duo) to do login
        return redirect(auth.login())
        # If AuthNRequest ID need to be stored in order to later validate it, do instead
        # sso_built_url = auth.login()
        # request.session['AuthNRequestID'] = auth.get_last_request_id()
        # return redirect(sso_built_url)
    elif "sso2" in request.args:
        return_to = "%sattrs/" % request.host_url
        return redirect(auth.login(return_to))
    elif "slo" in request.args:
        ## SCOTT: this is for logging out. 
        name_id = session_index = name_id_format = name_id_nq = name_id_spnq = None
        if "samlNameId" in session:
            name_id = session["samlNameId"]
        if "samlSessionIndex" in session:
            session_index = session["samlSessionIndex"]
        if "samlNameIdFormat" in session:
            name_id_format = session["samlNameIdFormat"]
        if "samlNameIdNameQualifier" in session:
            name_id_nq = session["samlNameIdNameQualifier"]
        if "samlNameIdSPNameQualifier" in session:
            name_id_spnq = session["samlNameIdSPNameQualifier"]

        return redirect(auth.logout(name_id=name_id, session_index=session_index, nq=name_id_nq, name_id_format=name_id_format, spnq=name_id_spnq))
    elif "acs" in request.args:
        ## SCOTT: This is where Duo redirects to after login
        sk = ','.join(session.keys())
        fk = ','.join(request.form.keys())
        app.logger.info(f'SCOTTSAML acs session keys: {sk} form keys: {fk}')
        request_id = None
        if "AuthNRequestID" in session:
            request_id = session["AuthNRequestID"]

        auth.process_response(request_id=request_id)
        errors = auth.get_errors()
        app.logger.info(f'SCOTTSAML errors: {errors}')
        not_auth_warn = not auth.is_authenticated()
        if len(errors) == 0:
            if "AuthNRequestID" in session:
                del session["AuthNRequestID"]
            session["samlUserdata"] = auth.get_attributes()
            session["samlNameId"] = auth.get_nameid()
            session["samlNameIdFormat"] = auth.get_nameid_format()
            session["samlNameIdNameQualifier"] = auth.get_nameid_nq()
            session["samlNameIdSPNameQualifier"] = auth.get_nameid_spnq()
            session["samlSessionIndex"] = auth.get_session_index()
            self_url = OneLogin_Saml2_Utils.get_self_url(req)
            if "RelayState" in request.form and self_url != request.form["RelayState"]:
                # To avoid 'Open Redirect' attacks, before execute the redirection confirm
                # the value of the request.form['RelayState'] is a trusted URL.
                url = auth.redirect_to(request.form["RelayState"])
                # return redirect(auth.redirect_to(request.form["RelayState"]))
                ## SCOTT: after login, let's redirect again
                return redirect(url_for('after_login'))
        elif auth.get_settings().is_debug_active():
            error_reason = auth.get_last_error_reason()
    elif "sls" in request.args:
        request_id = None
        if "LogoutRequestID" in session:
            request_id = session["LogoutRequestID"]
        dscb = lambda: session.clear()
        url = auth.process_slo(request_id=request_id, delete_session_cb=dscb)
        errors = auth.get_errors()
        if len(errors) == 0:
            if url is not None:
                # To avoid 'Open Redirect' attacks, before execute the redirection confirm
                # the value of the url is a trusted URL.
                return redirect(url)
            else:
                success_slo = True
        elif auth.get_settings().is_debug_active():
            error_reason = auth.get_last_error_reason()

    if "samlUserdata" in session:
        paint_logout = True
        if len(session["samlUserdata"]) > 0:
            attributes = session["samlUserdata"].items()

    return render_template("index.html", errors=errors, error_reason=error_reason, not_auth_warn=not_auth_warn, success_slo=success_slo, attributes=attributes, paint_logout=paint_logout)


@app.route("/attrs/")
def attrs():
    paint_logout = False
    attributes = False

    if "samlUserdata" in session:
        paint_logout = True
        if len(session["samlUserdata"]) > 0:
            attributes = session["samlUserdata"].items()

    return render_template("attrs.html", paint_logout=paint_logout, attributes=attributes)


@app.route("/metadata/")
def metadata():
    req = prepare_flask_request(request)
    auth = init_saml_auth(req)
    settings = auth.get_settings()
    metadata = settings.get_sp_metadata()
    errors = settings.validate_metadata(metadata)

    if len(errors) == 0:
        resp = make_response(metadata, 200)
        resp.headers["Content-Type"] = "text/xml"
    else:
        resp = make_response(", ".join(errors), 500)
    return resp

def multidict_to_html_str(dic):
    '''Returns HTML for a flask-style multi-dictionary'''
    val = 'multidict: \n<ul>'
    for key in dic.keys():
        val += f'    <li>{key}\n<ul>'
        for v in dic.getlist(key):
            val += f'    <li>{v}</li>'
        val += f'    </ul></li>'
    val += "</ul>"
    return val

def dic_to_html_str(dic):
    '''Returns HTML for a dictionary'''
    result = 'dictionary: \n<table><tr><th>key</th><th>value</th></tr>\n'
    for key,val in dic.items():
        result += f'<tr><td>{key}</td><td>{val}</td></tr>\n'
    result += "</table>\n"
    return result

@app.route("/home/")
def after_login():
    sess = 'session: '+dic_to_html_str(session)
    return f'Hello. You are logged in. Session is {sess}.'


if __name__ == "__main__":
    # the Cer and Key files are only readable by root, so root has to copy them
    context=('local.cer', 'local.key')
    app.run(host="0.0.0.0", port=1942, debug=True, ssl_context=context)