adding auth saml_user and starting work on login infrastructure

"""Provide authentication methods."""

from common.env import check_environment as ce

from typing import Dict, Any, Callable

from urllib.parse import urlparse, urlencode

import hashlib

import os

from typing import Callable, Any, Tuple

from common.mail import send_email

try:

    from onelogin.saml2.auth import OneLogin_Saml2_Auth

    from ldap3 import Server, Connection, ALL, SUBTREE

    from ldap3.core.exceptions import LDAPException, LDAPBindError

    from flask import jsonify

except ImportError:

    import sys

    from common.logz import create_logger

        return None

    connection.search(root_dn, user_search_filter, attributes=["*"])

    return connection.entries[0]

def authenticate_saml_user(request_scheme: str, request_host: str, request_url: str, request_method: str,

                      request_args: dict, request_form: dict,

                      success_handler: Callable = lambda _: None,

                      failure_handler: Callable = lambda _: None,

                      saml_path: str = None, logger=None):

    """

    Handles SAML authentication by processing the necessary components extracted from a web request.

    Parameters:

    - request_scheme (str): The request scheme ('http' or 'https').

    - request_host (str): The request host.

    - request_url (str): The full URL of the request.

    - request_method (str): The HTTP method of the request.

    - request_args (dict): The query parameters of the request as a dictionary.

    - request_form (dict): The form data of the request as a dictionary.

    - success_handler (Callable, optional): Function to handle successful SAML authentication.

      Defaults to a no-op function.

    - failure_handler (Callable, optional): Function to handle SAML authentication failure.

      Defaults to a no-op function.

    - saml_path (str, optional): Custom path for SAML configuration files. If not provided, environment variable

      'SAML_PATH' is used.

    - logger (Logger, optional): Logger object for logging. If not provided, no logging is performed.

    Returns:

    - Result of the `success_handler` or `failure_handler` callable depending on the SAML processing outcome.

    """

    url_data = urlparse(request_url)

    saml_request = {

        "https": "on" if request_scheme == "https" else "off",

        "http_host": request_host,

        "server_port": url_data.port,

        "script_name": url_data.path,

        "get_data": request_args,

        "post_data": request_form,

        "query_string": urlencode(request_args) if request_method == 'GET' else ''

    }

    saml_path = saml_path or ce("SAML_PATH")

    auth = OneLogin_Saml2_Auth(saml_request, custom_base_path=saml_path)

    try:

        if "sso" in request_args:

            return success_handler(auth.login())

        elif "slo" in request_args:

            return success_handler(auth.logout())

        elif "acs" in request_args:

            auth.process_response(request_id=None)  # Assuming request carries all needed info

            if auth.is_authenticated():

                return success_handler(None)  # pass user details / attributes

            else:

                return failure_handler("SAML authentication failed")

        else:

            return failure_handler("No valid SAML action found")

    except Exception as e:

        error_msg = f"SAML authentication error: {str(e)}"

        if logger:

            logger.error(error_msg)

        return failure_handler(error_msg)

def generate_salt() -> str:

    return str(os.urandom(32)).replace("\\", "").replace("b", "")

def hash_password(password: str, salt: str, hash_algorithm: Callable[..., Any] = hashlib.pbkdf2_hmac,

                  *args: Any, **kwargs: Any) -> str:

    return hash_algorithm(password.encode(), salt.encode(), *args, **kwargs).hex()

def login(username: str, password: str, db: str, table_name: str, username_column: str, password_column: str,

          salt_column: str) -> dict:

    user_data = db.query(f"SELECT {password_column}, {salt_column} FROM {table_name} WHERE {username_column} = '{username}'")

    if len(user_data) == 0:

        return {"msg": "Username does not exist."}, 401

    user_pass, salt = user_data[0]

    pass_to_check = hash_password(password, salt)

    if pass_to_check == user_pass:

        # User authentication successful

        return {"msg": "Login successful!"}

    else:

        return {"msg": "Invalid uid or password"}, 401

def register_user(username: str, password: str, confirm_pw: str, first_name: str, last_name: str, email: str,

                  work_sector: str, user_type: str, reason: str, db: UNDB, table_name: str, username_column: str,

                  password_column: str, salt_column: str) -> dict:

    # Check if username already exists

    if check_username_exists(username, db, table_name, username_column):

        return {"msg": "Username already exists."}, 401

    # Check if passwords match

    if password != confirm_pw:

        return {"msg": "Passwords do not match."}, 401

    # Generate salt and hash the password

    salt = generate_salt()

    hashed_pw = hash_password(password, salt)

    try:

        # Insert user data into the database

        db.cursor.execute(f"""INSERT INTO {table_name} ({username_column}, first_name, last_name, email, work_sector, 

                             user_type, {password_column}, {salt_column}, last_login_date, enabled) VALUES 

                             ('{username}', '{first_name}', '{last_name}', '{email}', '{work_sector}', 

                             '{user_type}', '{hashed_pw}', '{salt}', CURRENT_TIMESTAMP, false);""")

        db.commit()

    except Exception as e:

        return {"msg": "Error registering new user account.", "error": str(e)}, 401

    # Send email notification

    os.environ["EMAIL_RECIPIENTS"] = "plattmw@ornl.gov, burdetteja@ornl.gov"

    os.environ["EMAIL_SENDER"] = "psplatial@ornl.gov"

    subject = "PS Platial Account Request"

    msg = f"{last_name}, {first_name} at {email} has requested a PlanetSense Platial account for the following reason: {reason}."

    send_email(subject, msg)

    return {"msg": "Thank you for submitting an account request. Our team will review your request and respond as soon as possible."}

def reset_password(username: str, token: str, new_password: str, confirm_new_password: str, db: UNDB,

                   table_name: str, username_column: str, salt_column: str) -> dict:

    # Retrieve user's salt from the database

    user_salt = db.query(f"SELECT {salt_column} FROM {table_name} WHERE {username_column} = '{username}';")

    if not user_salt:

        return {"msg": "Username not found."}, 401

    user_salt = str(user_salt[0][0]).replace("b", "")

    # Validate token

    token_to_match = db.query(f"SELECT email_reset_code FROM {table_name} WHERE {username_column} = '{username}';")

    if token != token_to_match[0][0]:

        return {"msg": "Token does not match."}, 401

    # Check if passwords match

    if new_password != confirm_new_password:

        return {"msg": "Passwords do not match."}, 401

    # Generate new hashed password and update in the database

    new_hashed_pw = hash_password(new_password, user_salt)

    try:

        db.cursor.execute(f"UPDATE {table_name} SET {password_column} = '{new_hashed_pw}' WHERE {username_column} = '{username}';")

        db.commit()

    except Exception as e:

        return {"msg": "Failed to update password.", "error": str(e)}, 401

    return {"msg": "Password updated successfully! Redirecting you back to login page now..."}