# -*- coding: utf-8 -*-
# Copyright 2012-2021 CERN
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# Authors:
# - Mario Lassnig <mario.lassnig@cern.ch>, 2012-2021
# - Martin Barisits <martin.barisits@cern.ch>, 2012-2019
# - Vincent Garonne <vincent.garonne@cern.ch>, 2012-2017
# - Angelos Molfetas <Angelos.Molfetas@cern.ch>, 2012
# - Ralph Vigne <ralph.vigne@cern.ch>, 2013-2014
# - Cedric Serfon <cedric.serfon@cern.ch>, 2013
# - Thomas Beermann <thomas.beermann@cern.ch>, 2017
# - Hannes Hansen <hannes.jakob.hansen@cern.ch>, 2018-2019
# - Ruturaj Gujar <ruturaj.gujar23@gmail.com>, 2019
# - Jaroslav Guenther <jaroslav.guenther@cern.ch>, 2019-2020
# - Benedikt Ziemons <benedikt.ziemons@cern.ch>, 2020-2021
# - Rahul Chauhan <omrahulchauhan@gmail.com>, 2021
# - David Población Criado <david.poblacion.criado@cern.ch>, 2021

import datetime
import hashlib
import random
import sys
import traceback
from base64 import b64decode

import paramiko
from dogpile.cache import make_region
from dogpile.cache.api import NO_VALUE
from sqlalchemy import and_, or_, select, delete

from rucio.common.exception import CannotAuthenticate, RucioException
from rucio.common.utils import generate_uuid
from rucio.core.account import account_exists
from rucio.core.oidc import validate_jwt
from rucio.db.sqla import filter_thread_work
from rucio.db.sqla import models
from rucio.db.sqla.constants import IdentityType
from rucio.db.sqla.session import read_session, transactional_session


def token_key_generator(namespace, fni, **kwargs):
    """ :returns: generate key function """
    def generate_key(token, session=None):
        """ :returns: token """
        return token
    return generate_key


TOKENREGION = make_region(
    function_key_generator=token_key_generator
).configure(
    'dogpile.cache.memory',
    expiration_time=3600
)


@transactional_session
def get_auth_token_user_pass(account, username, password, appid, ip=None, session=None):
    """
    Authenticate a Rucio account temporarily via username and password.

    The token lifetime is 1 hour.

    :param account: Account identifier as a string.
    :param username: Username as a string.
    :param password: SHA1 hash of the password as a string.
    :param appid: The application identifier as a string.
    :param ip: IP address of the client a a string.
    :param session: The database session in use.

    :returns: A dict with token and expires_at entries.
    """

    # Make sure the account exists
    if not account_exists(account, session=session):
        return None

    result = session.query(models.Identity).filter_by(identity=username,
                                                      identity_type=IdentityType.USERPASS).first()

    db_salt = result['salt']
    db_password = result['password']

    salted_password = db_salt + password.encode()
    if db_password != hashlib.sha256(salted_password).hexdigest():
        return None

    # get account identifier
    result = session.query(models.IdentityAccountAssociation).filter_by(identity=username,
                                                                        identity_type=IdentityType.USERPASS,
                                                                        account=account).first()
    db_account = result['account']

    # remove expired tokens
    __delete_expired_tokens_account(account=account, session=session)

    # create new rucio-auth-token for account
    tuid = generate_uuid()  # NOQA
    token = '%(account)s-%(username)s-%(appid)s-%(tuid)s' % locals()
    new_token = models.Token(account=db_account, identity=username, token=token, ip=ip)
    new_token.save(session=session)

    return token_dictionary(new_token)


@transactional_session
def get_auth_token_x509(account, dn, appid, ip=None, session=None):
    """
    Authenticate a Rucio account temporarily via an x509 certificate.

    The token lifetime is 1 hour.

    :param account: Account identifier as a string.
    :param dn: Client certificate distinguished name string, as extracted by Apache/mod_ssl.
    :param id: The application identifier as a string.
    :param ipaddr: IP address of the client as a string.
    :param session: The database session in use.

    :returns: A dict with token and expires_at entries.
    """

    # Make sure the account exists
    if not account_exists(account, session=session):
        return None

    # remove expired tokens
    __delete_expired_tokens_account(account=account, session=session)

    # create new rucio-auth-token for account
    tuid = generate_uuid()  # NOQA
    token = '%(account)s-%(dn)s-%(appid)s-%(tuid)s' % locals()
    new_token = models.Token(account=account, identity=dn, token=token, ip=ip)
    new_token.save(session=session)

    return token_dictionary(new_token)


@transactional_session
def get_auth_token_gss(account, gsstoken, appid, ip=None, session=None):
    """
    Authenticate a Rucio account temporarily via a GSS token.

    The token lifetime is 1 hour.

    :param account: Account identifier as a string.
    :param gsscred: GSS principal@REALM as a string.
    :param appid: The application identifier as a string.
    :param ip: IP address of the client as a string.
    :param session: The database session in use.

    :returns: A dict with token and expires_at entries.
    """

    # Make sure the account exists
    if not account_exists(account, session=session):
        return None

    # remove expired tokens
    __delete_expired_tokens_account(account=account, session=session)

    # create new rucio-auth-token for account
    tuid = generate_uuid()  # NOQA
    token = '%(account)s-%(gsstoken)s-%(appid)s-%(tuid)s' % locals()
    new_token = models.Token(account=account, token=token, ip=ip)
    new_token.save(session=session)

    return token_dictionary(new_token)


@transactional_session
def get_auth_token_ssh(account, signature, appid, ip=None, session=None):
    """
    Authenticate a Rucio account temporarily via SSH key exchange.

    The token lifetime is 1 hour.

    :param account: Account identifier as a string.
    :param signature: Response to server challenge signed with SSH private key as string.
    :param appid: The application identifier as a string.
    :param ip: IP address of the client as a string.
    :param session: The database session in use.

    :returns: A dict with token and expires_at entries.
    """
    if not isinstance(signature, bytes):
        signature = signature.encode()

    # Make sure the account exists
    if not account_exists(account, session=session):
        return None

    # get all active challenge tokens for the requested account
    active_challenge_tokens = session.query(models.Token).filter(models.Token.expired_at >= datetime.datetime.utcnow(),
                                                                 models.Token.account == account,
                                                                 models.Token.token.like('challenge-%')).all()

    # get all identities for the requested account
    identities = session.query(models.IdentityAccountAssociation).filter_by(identity_type=IdentityType.SSH,
                                                                            account=account).all()

    # no challenge tokens found
    if not active_challenge_tokens:
        return None

    # try all available SSH identities for the account with the provided signature
    match = False
    for identity in identities:
        data = identity['identity'].split()[1]
        data += '=' * ((4 - len(data) % 4) % 4)  # adding required padding
        pub_k = paramiko.RSAKey(data=b64decode(data))
        for challenge_token in active_challenge_tokens:
            if pub_k.verify_ssh_sig(str(challenge_token['token']).encode(),
                                    paramiko.Message(signature)):
                match = True
                break
        if match:
            break

    if not match:
        return None

    # remove expired tokens
    __delete_expired_tokens_account(account=account, session=session)

    # create new rucio-auth-token for account
    tuid = generate_uuid()  # NOQA
    token = '%(account)s-ssh:pubkey-%(appid)s-%(tuid)s' % locals()
    new_token = models.Token(account=account, token=token, ip=ip)
    new_token.save(session=session)

    return token_dictionary(new_token)


@transactional_session
def get_ssh_challenge_token(account, appid, ip=None, session=None):
    """
    Prepare a challenge token for subsequent SSH public key authentication.

    The challenge lifetime is fixed to 10 seconds.

    :param account: Account identifier as a string.
    :param appid: The application identifier as a string.
    :param ip: IP address of the client as a string.

    :returns: A dict with token and expires_at entries.
    """

    # Make sure the account exists
    if not account_exists(account, session=session):
        return None

    # Cryptographically secure random number.
    # This requires a /dev/urandom like device from the OS
    rng = random.SystemRandom()
    crypto_rand = rng.randint(0, sys.maxsize)

    # give the client 10 seconds max to sign the challenge token
    expiration = datetime.datetime.utcnow() + datetime.timedelta(seconds=10)
    expiration_unix = expiration.strftime("%s")

    challenge_token = 'challenge-%(crypto_rand)s-%(account)s-%(expiration_unix)s' % locals()

    new_challenge_token = models.Token(account=account, token=challenge_token, ip=ip,
                                       expired_at=expiration)
    new_challenge_token.save(session=session)

    return token_dictionary(new_challenge_token)


@transactional_session
def get_auth_token_saml(account, saml_nameid, appid, ip=None, session=None):
    """
    Authenticate a Rucio account temporarily via SAML.

    The token lifetime is 1 hour.

    :param account: Account identifier as a string.
    :param saml_nameid: SAML NameID of the client.
    :param appid: The application identifier as a string.
    :param ip: IP address of the client a a string.
    :param session: The database session in use.

    :returns: A dict with token and expires_at entries.
    """

    # Make sure the account exists
    if not account_exists(account, session=session):
        return None

    # remove expired tokens
    __delete_expired_tokens_account(account=account, session=session)

    tuid = generate_uuid()  # NOQA
    token = '%(account)s-%(saml_nameid)s-%(appid)s-%(tuid)s' % locals()
    new_token = models.Token(account=account, identity=saml_nameid, token=token, ip=ip)
    new_token.save(session=session)

    return token_dictionary(new_token)


@transactional_session
def redirect_auth_oidc(auth_code, fetchtoken=False, session=None):
    """
    Finds the Authentication URL in the Rucio DB oauth_requests table
    and redirects user's browser to this URL.

    :param auth_code: Rucio assigned code to redirect
                      authorization securely to IdP via Rucio Auth server through a browser.
    :param fetchtoken: If True, valid token temporarily saved in the oauth_requests table
                       will be returned. If False, redirection URL is returned.
    :param session: The database session in use.

    :returns: result of the query (authorization URL or a
              token if a user asks with the correct code) or None.
              Exception thrown in case of an unexpected crash.

    """
    try:
        redirect_result = session.query(models.OAuthRequest.redirect_msg).filter_by(access_msg=auth_code).first()

        if not redirect_result:
            return None

        if 'http' not in redirect_result[0] and fetchtoken:
            # in this case the function check if the value is a valid token
            vdict = validate_auth_token(redirect_result[0], session=session)
            if vdict:
                return redirect_result[0]
            return None
        elif 'http' in redirect_result[0] and not fetchtoken:
            # return redirection URL
            return redirect_result[0]
    except:
        raise CannotAuthenticate(traceback.format_exc())


@transactional_session
def delete_expired_tokens(total_workers, worker_number, limit=1000, session=None):
    """
    Delete expired tokens.

    :param total_workers:      Number of total workers.
    :param worker_number:      id of the executing worker.
    :param limit:              Maximum number of tokens to delete.
    :param session:            Database session in use.

    :returns: number of deleted rows
    """

    # get expired tokens
    try:
        # delete all expired tokens except tokens which have refresh token that is still valid
        query = session.query(models.Token.token).filter(and_(models.Token.expired_at <= datetime.datetime.utcnow()))\
                                                 .filter(or_(models.Token.refresh_expired_at.__eq__(None),
                                                             models.Token.refresh_expired_at <= datetime.datetime.utcnow()))\
                                                 .order_by(models.Token.expired_at)

        query = filter_thread_work(session=session, query=query, total_threads=total_workers, thread_id=worker_number, hash_variable='token')

        # limiting the number of tokens deleted at once
        query = query.limit(limit)
        # remove expired tokens
        deleted_tokens = 0
        for items in session.execute(query).partitions(10):
            tokens = tuple(map(lambda row: row.token, items))
            deleted_tokens += session.query(models.Token) \
                                     .filter(models.Token.token.in_(tokens)) \
                                     .with_for_update(skip_locked=True) \
                                     .delete(synchronize_session='fetch')

    except Exception as error:
        raise RucioException(error.args)

    return deleted_tokens


@read_session
def query_token(token, session=None):
    """
    Validate an authentication token using the database. This method will only be called
    if no entry could be found in the according cache.

    :param token: Authentication token as a variable-length string.
    :param session: The database session in use.

    :returns: dictionary { account: <account name>,
                           identity: <identity>,
                           lifetime: <token lifetime>,
                           audience: <audience>,
                           authz_scope: <authz_scope> }
              if successful, None otherwise.
    """
    # Query the DB to validate token
    ret = session.query(models.Token.account,
                        models.Token.identity,
                        models.Token.expired_at,
                        models.Token.audience,
                        models.Token.oidc_scope).\
        filter(models.Token.token == token,
               models.Token.expired_at > datetime.datetime.utcnow()).\
        all()
    if ret:
        return {'account': ret[0][0],
                'identity': ret[0][1],
                'lifetime': ret[0][2],
                'audience': ret[0][3],
                'authz_scope': ret[0][4]}
    return None


@transactional_session
def validate_auth_token(token, session=None):
    """
    Validate an authentication token.

    :param token: Authentication token as a variable-length string.

    :returns: dictionary { account: <account name>,
                           identity: <identity>,
                           lifetime: <token lifetime>,
                           audience: <audience>,
                           authz_scope: <authz_scope> }
              if successful, None otherwise.
    """
    if not token:
        return None

    # Be gentle with bash variables, there can be whitespace
    token = token.strip()

    # Check if token ca be found in cache region
    value = TOKENREGION.get(token)
    if value is NO_VALUE:  # no cached entry found
        value = query_token(token, session=session)
        if not value:
            # identify JWT access token and validte
            # & save it in Rucio if scope and audience are correct
            if len(token.split(".")) == 3:
                value = validate_jwt(token, session=session)
                if not value:
                    return None
            else:
                return None
        # save token in the cache
        TOKENREGION.set(token, value)
    if value.get('lifetime', datetime.datetime(1970, 1, 1)) < datetime.datetime.utcnow():  # check if expired
        TOKENREGION.delete(token)
        return None
    return value


def token_dictionary(token: models.Token):
    return {'token': token.token, 'expires_at': token.expired_at}


@transactional_session
def __delete_expired_tokens_account(account, session=None):
    """"
    Deletes expired tokens from the database.

    :param account: Account to delete expired tokens.
    :param session: The database session in use.
    """
    stmt_select = select(models.Token) \
        .where(and_(models.Token.expired_at < datetime.datetime.utcnow(),
                    models.Token.account == account)) \
        .with_for_update(skip_locked=True)
    result_select = session.execute(stmt_select)

    tokens = []
    for t in result_select.columns('token'):
        tokens.append(t.token)

    stmt_delete = delete(models.Token) \
        .where(models.Token.token.in_(tokens)) \
        .prefix_with("/*+ INDEX(TOKENS_ACCOUNT_EXPIRED_AT_IDX) */")
    session.execute(stmt_delete)
