/*
 * Note: Information concerning change history is maintained by the version control system. This information, along with code relationship (e.g. used by, uses,
 * dependencies), is readily available through the Eclipse Integrated Development Environment.
 */
package com.lmco.spacefence.server.certificaterevocation;

import org.bouncycastle.asn1.ASN1InputStream;
import org.bouncycastle.asn1.ASN1Primitive;
import org.bouncycastle.asn1.DERIA5String;
import org.bouncycastle.asn1.DEROctetString;
import org.bouncycastle.asn1.x509.CRLDistPoint;
import org.bouncycastle.asn1.x509.DistributionPoint;
import org.bouncycastle.asn1.x509.DistributionPointName;
import org.bouncycastle.asn1.x509.GeneralName;
import org.bouncycastle.asn1.x509.GeneralNames;

import java.io.ByteArrayInputStream;
import java.io.IOException;
import java.io.InputStream;
import java.net.URL;
import java.security.cert.CRLException;
import java.security.cert.Certificate;
import java.security.cert.CertificateException;
import java.security.cert.CertificateFactory;
import java.security.cert.CertificateParsingException;
import java.security.cert.X509CRL;
import java.util.ArrayList;
import java.util.Calendar;
import java.util.Collections;
import java.util.Date;
import java.util.List;
import java.util.TimeZone;
import java.util.concurrent.ConcurrentHashMap;

import io.undertow.UndertowLogger;
import io.undertow.server.HttpHandler;
import io.undertow.server.HttpServerExchange;
import io.undertow.server.SSLSessionInfo;
import io.undertow.server.ServerConnection;
import io.undertow.util.StatusCodes;

public class CertificateRevocationFilter implements HttpHandler {

    private static final String CRLDistributionPointsOID = "2.5.29.31";

    private HttpHandler next;
    private final ConcurrentHashMap<String, X509CRL> crlMap = new ConcurrentHashMap<>();

    public CertificateRevocationFilter(HttpHandler next) {
        this.next = next;
    }

    @Override
    public void handleRequest(HttpServerExchange exchange) throws Exception {
        ServerConnection serverConnection = exchange.getConnection();
        SSLSessionInfo sslSessionInfo = serverConnection.getSslSessionInfo();
        if (sslSessionInfo != null) {
            Certificate[] chain = sslSessionInfo.getPeerCertificates();
            if (chain != null) {
                Date now = Calendar.getInstance(TimeZone.getTimeZone("UTC")).getTime();

                for (Certificate cert : chain) {
                    if (cert instanceof java.security.cert.X509Certificate) {
                        java.security.cert.X509Certificate x509Cert = (java.security.cert.X509Certificate) cert;
                        List<String> crlDistributionPoints = getCrlDistributionPoints(x509Cert);
                        for (String crlDistributionPoint : crlDistributionPoints) {
                            X509CRL crl = crlMap.get(crlDistributionPoint);
                            if (crl == null || now.after(crl.getNextUpdate())) {
                                crl = downloadCRLFromWeb(crlDistributionPoint);
                                if (crl == null) {
                                    UndertowLogger.REQUEST_LOGGER.info("Unable to download CRL for Certificate SN:" + x509Cert.getSerialNumber().toString(16) +
                                                                       " DN:" + x509Cert.getSubjectDN().getName() + " from URL:" + crlDistributionPoint);
                                    exchange.setStatusCode(StatusCodes.FORBIDDEN);
                                    exchange.endExchange();
                                    return;
                                } else {
                                    crlMap.put(crlDistributionPoint, crl);
                                }
                            }
                            if (crl.isRevoked(cert)) {
                                UndertowLogger.REQUEST_LOGGER.info("Revoked Certificate SN:" + x509Cert.getSerialNumber().toString(16) +
                                                                   " DN:" + x509Cert.getSubjectDN().getName());
                                exchange.setStatusCode(StatusCodes.FORBIDDEN);
                                exchange.endExchange();
                                return;
                            }
                        }
                    }
                }
                next.handleRequest(exchange);
            }
        } else {
            next.handleRequest(exchange);
        }
    }

    /**
     * Extracts all CRL distribution point URLs from the "CRL Distribution Point"
     * extension in a X.509 certificate. If CRL distribution point extension is
     * unavailable, returns an empty list.
     */
    private static List<String> getCrlDistributionPoints(java.security.cert.X509Certificate x509Cert) throws CertificateParsingException, IOException {
        byte[] crldpExt = x509Cert.getExtensionValue(CRLDistributionPointsOID);
        if (crldpExt == null) {
            return Collections.emptyList();
        }

        ASN1InputStream oAsnInStream = new ASN1InputStream(new ByteArrayInputStream(crldpExt));
        ASN1Primitive derObjCrlDP = oAsnInStream.readObject();
        DEROctetString dosCrlDP = (DEROctetString) derObjCrlDP;
        byte[] crldpExtOctets = dosCrlDP.getOctets();
        ASN1InputStream oAsnInStream2 = new ASN1InputStream(new ByteArrayInputStream(crldpExtOctets));
        ASN1Primitive derObj2 = oAsnInStream2.readObject();
        CRLDistPoint distPoint = CRLDistPoint.getInstance(derObj2);
        List<String> crlUrls = new ArrayList<>();
        for (DistributionPoint dp : distPoint.getDistributionPoints()) {
            DistributionPointName dpn = dp.getDistributionPoint();
            // Look for URIs in fullName
            if (dpn != null) {
                if (DistributionPointName.FULL_NAME == dpn.getType()) {
                    GeneralName[] genNames = GeneralNames.getInstance(dpn.getName()).getNames();
                    // Look for an URI
                    for (GeneralName genName : genNames) {
                        if (GeneralName.uniformResourceIdentifier == genName.getTagNo()) {
                            String url = DERIA5String.getInstance(genName.getName()).getString();
                            crlUrls.add(url);
                        }
                    }
                }
            }
        }
        return crlUrls;
    }

    /**
     * Downloads a CRL from given HTTP/HTTPS/FTP URL, e.g.
     * http://crl.infonotary.com/crl/identity-ca.crl
     */
    private static X509CRL downloadCRLFromWeb(String crlURL) throws IOException, CertificateException, CRLException {
        try (InputStream crlStream = new URL(crlURL).openStream()) {
            CertificateFactory cf = CertificateFactory.getInstance("X.509");
            return (X509CRL) cf.generateCRL(crlStream);
        }
    }
}