/*
 * Copyright (c) 2016, WSO2 Inc. (http://www.wso2.org) All Rights Reserved.
 *
 * WSO2 Inc. licenses this file to you 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.
 */
package org.wso2.carbon.apimgt.handlers;

import com.google.gson.Gson;
import org.apache.axis2.context.MessageContext;
import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
import org.apache.synapse.core.axis2.Axis2MessageContext;
import org.apache.synapse.rest.AbstractHandler;
import org.wso2.carbon.apimgt.handlers.beans.Certificate;
import org.wso2.carbon.apimgt.handlers.beans.ValidationResponce;
import org.wso2.carbon.apimgt.handlers.config.IOTServerConfiguration;
import org.wso2.carbon.apimgt.handlers.invoker.RESTInvoker;
import org.wso2.carbon.apimgt.handlers.invoker.RESTResponse;
import org.wso2.carbon.apimgt.handlers.utils.AuthConstants;
import org.wso2.carbon.apimgt.handlers.utils.Utils;
import io.entgra.device.mgt.core.certificate.mgt.core.dto.CertificateResponse;
import io.entgra.device.mgt.core.certificate.mgt.core.exception.KeystoreException;
import io.entgra.device.mgt.core.certificate.mgt.core.impl.CertificateGenerator;
import org.wso2.carbon.context.PrivilegedCarbonContext;

import javax.security.cert.CertificateEncodingException;
import java.io.ByteArrayInputStream;
import java.io.IOException;
import java.net.URI;
import java.net.URISyntaxException;
import java.security.cert.CertificateException;
import java.security.cert.CertificateFactory;
import java.security.cert.X509Certificate;
import java.util.HashMap;
import java.util.Map;
import java.util.StringTokenizer;

/**
 * Synapse gateway handler for API authentication.
 */
public class AuthenticationHandler extends AbstractHandler {
    private static final Log log = LogFactory.getLog(AuthenticationHandler.class);
    private RESTInvoker restInvoker;

    private static final String X_JWT_ASSERTION = "X-JWT-Assertion";
    private static final String JWTTOKEN = "JWTToken";
    private static final String AUTHORIZATION = "Authorization";
    private static final String BEARER = "Basic ";
    private static final String CONTENT_TYPE = "Content-Type";
    private static final boolean USE_INTERNAL_CERT_VERIFIER = true;

    private IOTServerConfiguration iotServerConfiguration;

    /**
     * Setting up configurations at the constructor
     */
    public AuthenticationHandler() {
        log.info("Engaging API Security Handler..........");
        restInvoker = new RESTInvoker();
        this.iotServerConfiguration = Utils.initConfig();
    }

    /**
     * Handling the message and checking the security.
     *
     * @param messageContext Request message context.
     * @return Boolean value of the result of the processing the request.
     */
    @Override
    public boolean handleRequest(org.apache.synapse.MessageContext messageContext) {
        org.apache.axis2.context.MessageContext axisMC = ((Axis2MessageContext) messageContext).getAxis2MessageContext();

        String ctxPath = messageContext.getTo().getAddress().trim();

        if (log.isDebugEnabled()) {
            log.debug("Authentication handler invoked by: " + ctxPath);
        }
        Map<String, String> headers = (Map<String, String>) axisMC.getProperty(MessageContext.TRANSPORT_HEADERS);
        try {
            int tenantId = PrivilegedCarbonContext.getThreadLocalCarbonContext().getTenantId();
            RESTResponse response = null;
            if (headers.containsKey(AuthConstants.MDM_SIGNATURE)) {

                String mdmSignature = headers.get(AuthConstants.MDM_SIGNATURE);
                if (log.isDebugEnabled()) {
                    log.debug("Verify Cert:\n" + mdmSignature);
                }
                URI certVerifyUrl = new URI(iotServerConfiguration.getVerificationEndpoint() + "ios");
                Map<String, String> certVerifyHeaders = this.setHeaders();

                Certificate certificate = new Certificate();
                certificate.setPem(mdmSignature);
                certificate.setTenantId(tenantId);
                certificate.setSerial("");

                Gson gson = new Gson();
                String certVerifyContent = gson.toJson(certificate);
                response = restInvoker.invokePOST(certVerifyUrl, certVerifyHeaders, certVerifyContent);

                String str = response.getContent();
                if (log.isDebugEnabled()) {
                    log.debug("Verify response:" + response.getContent());
                    log.debug("Response String : " + str);
                }
                if (response.getHttpStatus() == 200 && str.contains(JWTTOKEN)) {
                    ValidationResponce validationResponce = gson.fromJson(str, ValidationResponce.class);
                    headers.put(X_JWT_ASSERTION, validationResponce.getJWTToken());
                } else {
                    return false;
                }

            } else if (headers.containsKey(AuthConstants.PROXY_MUTUAL_AUTH_HEADER)) {
                String subjectDN = headers.get(AuthConstants.PROXY_MUTUAL_AUTH_HEADER);

                if (log.isDebugEnabled()) {
                    log.debug("Verify subject DN: " + subjectDN);
                }

                if (USE_INTERNAL_CERT_VERIFIER) {
                    CertificateResponse certificateResponse = Utils.getCertificateManagementService()
                            .verifySubjectDN(subjectDN);
                    if (certificateResponse != null && certificateResponse.getCommonName() != null
                            && !certificateResponse.getCommonName().isEmpty()) {
                        return true;
                    }
                } else {
                    String deviceType = this.getDeviceType(messageContext.getTo().getAddress().trim());
                    URI certVerifyUrl = new URI(iotServerConfiguration.getVerificationEndpoint() + deviceType);
                    Map<String, String> certVerifyHeaders = this.setHeaders();

                    Certificate certificate = new Certificate();
                    certificate.setPem(subjectDN);
                    certificate.setTenantId(tenantId);
                    certificate.setSerial(AuthConstants.PROXY_MUTUAL_AUTH_HEADER);

                    Gson gson = new Gson();
                    String certVerifyContent = gson.toJson(certificate);
                    response = restInvoker.invokePOST(certVerifyUrl, certVerifyHeaders, certVerifyContent);
                    if (log.isDebugEnabled()) {
                        log.debug("Verify response:" + response.getContent());
                    }
                }
            } else if (headers.containsKey(AuthConstants.MUTUAL_AUTH_HEADER)) {
                javax.security.cert.X509Certificate[] certs =
                        (javax.security.cert.X509Certificate[]) axisMC.getProperty(AuthConstants.CLIENT_CERTIFICATE);
                CertificateFactory cf = CertificateFactory.getInstance("X.509");
                ByteArrayInputStream bais = new ByteArrayInputStream(certs[0].getEncoded());
                X509Certificate x509 = (X509Certificate) cf.generateCertificate(bais);
                bais.close();
                if (x509 != null) {
                    headers.put(AuthConstants.PROXY_MUTUAL_AUTH_HEADER, CertificateGenerator.getCommonName(x509));
                    return true;
                }
            } else if (headers.containsKey(AuthConstants.ENCODED_PEM)) {
                String encodedPem = headers.get(AuthConstants.ENCODED_PEM);
                if (log.isDebugEnabled()) {
                    log.debug("Verify Cert:\n" + encodedPem);
                }
                String deviceType = this.getDeviceType(messageContext.getTo().getAddress().trim());
                URI certVerifyUrl = new URI(iotServerConfiguration.getVerificationEndpoint() + deviceType);
                Map<String, String> certVerifyHeaders = this.setHeaders();

                Certificate certificate = new Certificate();
                certificate.setPem(encodedPem);
                certificate.setTenantId(tenantId);
                certificate.setSerial("");
                Gson gson = new Gson();
                String certVerifyContent = gson.toJson(certificate);
                response = restInvoker.invokePOST(certVerifyUrl, certVerifyHeaders, certVerifyContent);
                if (log.isDebugEnabled()) {
                    log.debug("Verify response:" + response.getContent());
                }
            } else {
                log.warn("Unauthorized request for api: " + ctxPath);
                return false;
            }
            if (response != null && !response.getContent().contains("invalid")) {
                return true;
            }
            log.warn("Unauthorized request for api: " + ctxPath);
            return false;
        } catch (IOException e) {
            log.error("Error while processing certificate.", e);
            return false;
        } catch (URISyntaxException e) {
            log.error("Error while processing certificate.", e);
            return false;
        } catch (CertificateException e) {
            log.error("Certificate issue occurred when generating converting PEM to x509Certificate", e);
            return false;
        } catch (CertificateEncodingException e) {
            log.error("Error while attempting to encode certificate.", e);
            return false;
        } catch (KeystoreException e) {
            log.error("Error while attempting to validate certificate.", e);
            return false;
        }
    }

    @Override
    public boolean handleResponse(org.apache.synapse.MessageContext messageContext) {
        return true;
    }


    private String getDeviceType(String url) {
        StringTokenizer parts = new StringTokenizer(url, "/");
        while (parts.hasMoreElements()) {
            if (parts.nextElement().equals("device-mgt")) {
                return (String) parts.nextElement();
            }
        }
        return null;
    }

    private Map<String, String> setHeaders() {
        Map<String, String> map = new HashMap<>();
        String accessToken = Utils.getBase64EncodedToken(iotServerConfiguration);
        map.put(AUTHORIZATION, BEARER + accessToken);
        map.put(CONTENT_TYPE, "application/json");
        return map;
    }
}
