package com.hfg.cert;

import java.net.URL;
import java.security.KeyStore;
import java.security.PublicKey;
import java.security.cert.*;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collection;
import java.util.Collections;
import java.util.HashSet;
import java.util.List;
import java.util.Set;
import java.util.regex.Matcher;
import java.util.regex.Pattern;
import java.util.stream.Collectors;
import javax.net.ssl.HttpsURLConnection;
import javax.net.ssl.SSLContext;
import javax.net.ssl.SSLSocket;
import javax.net.ssl.SSLSocketFactory;
import javax.net.ssl.TrustManager;
import javax.net.ssl.TrustManagerFactory;
import javax.net.ssl.X509TrustManager;

import com.hfg.util.StringBuilderPlus;
import com.hfg.util.StringUtil;
import com.hfg.util.collection.OrderedSet;

//------------------------------------------------------------------------------
/**
 * General certificate utility methods.
 *
 * @author J. Alex Taylor, hairyfatguy.com
 */
//------------------------------------------------------------------------------
// com.hfg XML/HTML Coding Library
//
// This library is free software; you can redistribute it and/or
// modify it under the terms of the GNU Lesser General Public
// License as published by the Free Software Foundation; either
// version 2.1 of the License, or (at your option) any later version.
//
// This library is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
// Lesser General Public License for more details.
//
// You should have received a copy of the GNU Lesser General Public
// License along with this library; if not, write to the Free Software
// Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
//
// J. Alex Taylor, President, Founder, CEO, COO, CFO, OOPS hairyfatguy.com
// jataylor@hairyfatguy.com
//------------------------------------------------------------------------------

public class CertificateUtil
{

   //--------------------------------------------------------------------------
   public static void verifyCertificates(String inURL)
         throws Exception
   {
      Set<X509Certificate> certs;

      if (inURL.startsWith("https:"))
      {
         certs = getWebCertificates(inURL);
      }
      else if (inURL.startsWith("ldaps:"))
      {
         certs = geLdapCertificates(inURL);
      }
      else
      {
         throw new RuntimeException("Unsupported url protocol: " + StringUtil.singleQuote(inURL) + "!");
      }


      List<X509Certificate> trustedCerts = getTrustedCertificates();
      Set<X509Certificate> rootCAs = new HashSet<>(trustedCerts.size());
      Set<X509Certificate> intermediateCerts = new HashSet<>(trustedCerts.size());
      for (X509Certificate cert : trustedCerts)
      {
         if (isSelfSigned(cert))
         {
            rootCAs.add(cert);
         }
         else
         {
            intermediateCerts.add(cert);
         }
      }

      List<X509Certificate> reorderedCerts = new ArrayList<>(certs);
      if (reorderedCerts.size() > 1
          && reorderedCerts.get(0).getIssuerDN().equals(reorderedCerts.get(1).getSubjectDN()))
      {
         Collections.reverse(reorderedCerts);
      }

      for (X509Certificate cert : reorderedCerts)
      {
         System.out.println(generateCertSummary(cert));

         X509CertSelector selector = new X509CertSelector();
         selector.setCertificate(cert);

         Set<TrustAnchor> trustAnchors = new HashSet<>();
         for (X509Certificate rootCA : rootCAs)
         {
            trustAnchors.add(new TrustAnchor(rootCA, null));
         }

         PKIXBuilderParameters pkixParams = new PKIXBuilderParameters(trustAnchors, selector);

         pkixParams.setRevocationEnabled(false);

         CertStore intermediateCertStore = CertStore.getInstance("Collection",
                                                                 new CollectionCertStoreParameters(intermediateCerts));
         pkixParams.addCertStore(intermediateCertStore);

         CertPathBuilder builder = CertPathBuilder.getInstance("PKIX");
         PKIXCertPathBuilderResult result = (PKIXCertPathBuilderResult) builder.build(pkixParams);

         if (isSelfSigned(cert))
         {
            rootCAs.add(cert);
         }
         else
         {
            intermediateCerts.add(cert);
         }
      }
   }

   //--------------------------------------------------------------------------
   public static String generateCertSummary(X509Certificate inCert)
      throws CertificateParsingException
   {
      List<SAN> sanNames = getSubjectAlternateNames(inCert);

      StringBuilderPlus buffer = new StringBuilderPlus()
            .appendln("Certificate Info:")
            .appendln("----------------")
            .appendln("   Subject DN: " + inCert.getSubjectDN())
            .appendln("         Type: " + inCert.getType())
            .appendln("   Public Key Algorithm: " + inCert.getPublicKey().getAlgorithm())
            .appendln("   Public Key Format: " + inCert.getPublicKey().getFormat())
            .appendln(" Subject Alt Name(s): " + (sanNames != null ? StringUtil.join(sanNames, ", ") : ""))
            .appendln("    Issuer DN: " + inCert.getIssuerDN())
            .appendln("      Expires: " + inCert.getNotAfter());

      return buffer.toString();
   }

   //--------------------------------------------------------------------------
   public static Set<X509Certificate> getWebCertificates(String inURL)
         throws Exception
   {
      HttpsURLConnection conn = (HttpsURLConnection) new URL(inURL).openConnection();
      conn.connect();

      Set<X509Certificate> certSet = new OrderedSet<>(5);

      Certificate[] certs = conn.getServerCertificates();
      if (certs != null)
      {
         certSet.addAll((List<X509Certificate>) (Object) Arrays.asList(certs));
      }

      return certSet;
   }

   //--------------------------------------------------------------------------
   public static Set<X509Certificate> geLdapCertificates(String inURL)
         throws Exception
   {
      Pattern urlPattern = Pattern.compile("ldap(?:s)?://([\\w\\-\\.]+)(?:\\:(\\d+))?");
      Matcher m = urlPattern.matcher(inURL);
      if (! m.matches())
      {
         throw new RuntimeException("Unrecognized URL pattern: " + inURL);
      }

      String host = m.group(1);
      int port = (m.group(2) != null ? Integer.parseInt(m.group(2)) : 636);

      TrustManagerFactory trustMgrfactory = getTrustManagerFactory();

      X509TrustManager defaultTrustManager = (X509TrustManager) trustMgrfactory.getTrustManagers()[0];

      SavingTrustManager tm = new SavingTrustManager(defaultTrustManager);

      SSLSocket sslSocket = null;

      SSLContext context = SSLContext.getInstance("TLS");
      context.init(null, new TrustManager[]{tm}, null);
      SSLSocketFactory sslFactory = context.getSocketFactory();

      try
      {
         sslSocket = (SSLSocket) sslFactory.createSocket(host, port);
         sslSocket.setSoTimeout(5000);
         System.out.println("Starting SSL handshake...");
         sslSocket.startHandshake();
      }
      catch (Exception e)
      {
         // Ignore. We want to be able to examine the certificates even if there was
         // an error thrown because they aren't recognized or valid.
      }
      finally
      {
         if (sslSocket != null)
         {
            sslSocket.close();
         }
      }

      return tm.mNewCertificates;
   }

   //--------------------------------------------------------------------------
   private static List<X509Certificate> getTrustedCertificates()
      throws Exception
   {
      TrustManagerFactory factory = getTrustManagerFactory();

      List<TrustManager> trustManagers = Arrays.asList(factory.getTrustManagers());
      List<X509Certificate> certificates = trustManagers.stream()
            .filter(X509TrustManager.class::isInstance)
            .map(X509TrustManager.class::cast)
            .map(trustManager -> Arrays.asList(trustManager.getAcceptedIssuers()))
            .flatMap(Collection::stream)
            .collect(Collectors.toList());

      return certificates;
   }

   //--------------------------------------------------------------------------
   private static TrustManagerFactory getTrustManagerFactory()
      throws Exception
   {
      TrustManagerFactory factory = TrustManagerFactory.getInstance(TrustManagerFactory.getDefaultAlgorithm());
      factory.init((KeyStore) null);

      return factory;
   }

   //--------------------------------------------------------------------------
   public static boolean isSelfSigned(X509Certificate inCert)
   {
      boolean result = true;

      try
      {
         PublicKey key = inCert.getPublicKey();
         inCert.verify(key);
      }
      catch (Exception e)
      {
         result = false;
      }

      return result;
   }

   //--------------------------------------------------------------------------
   public static List<SAN> getSubjectAlternateNames(X509Certificate inCertificate)
      throws CertificateParsingException
   {
      List<SAN> values = null;
      try
      {
         /*
          The ASN.1 definition of the SubjectAltName extension is:
          SubjectAltName ::= GeneralNames

          GeneralNames :: = SEQUENCE SIZE (1..MAX) OF GeneralName

          GeneralName ::= CHOICE {
               otherName                       [0]     OtherName,
               rfc822Name                      [1]     IA5String,
               dNSName                         [2]     IA5String,
               x400Address                     [3]     ORAddress,
               directoryName                   [4]     Name,
               ediPartyName                    [5]     EDIPartyName,
               uniformResourceIdentifier       [6]     IA5String,
               iPAddress                       [7]     OCTET STRING,
               registeredID                    [8]     OBJECT IDENTIFIER}

          If this certificate does not contain a SubjectAltName extension, null is returned.
          Otherwise, a Collection is returned with an entry representing each GeneralName
          included in the extension. Each entry is a List whose first entry is an Integer
          (the name type, 0-8) and whose second entry is a String or a byte array (the name,
          in string or ASN.1 DER encoded form, respectively).
          */
         Collection<List<?>> altNames = inCertificate.getSubjectAlternativeNames();
         if (altNames != null)
         {
            values = new ArrayList<>(3);
            for (List item : altNames)
            {
               SAN_Type type = SAN_Type.valueOf((Integer) item.get(0));
               Object value = item.get(1);
               String stringValue = null;

               if (SAN_Type.OTHER_NAME.equals(type))
               {
                  if (value instanceof String)
                  {
                     stringValue = (String) value;
                  }
                  else if (value instanceof byte[])
                  {
                     // ASN.1 DER-encoded value
                     // TODO: Decode the value
                     stringValue = "DER-encoded value";
                        /*
                        try
                        {
                           ASN1InputStream decoder=null;
                           if(item.toArray()[1] instanceof byte[])
                              decoder = new ASN1InputStream((byte[]) item.toArray()[1]);
                           else if(item.toArray()[1] instanceof String)
                              identities.add( (String) item.toArray()[1] );
                           if(decoder==null) continue;
                           DEREncodable encoded = decoder.readObject();
                           encoded = ((DERSequence) encoded).getObjectAt(1);
                           encoded = ((DERTaggedObject) encoded).getObject();
                           encoded = ((DERTaggedObject) encoded).getObject();
                           String identity = ((DERUTF8String) encoded).getString();
                           identities.add(identity);
                        }
                        catch (UnsupportedEncodingException e) {
                           log.error("Error decoding subjectAltName" + e.getLocalizedMessage(),e);
                        }
                        catch (Exception e) {
                           log.error("Error decoding subjectAltName" + e.getLocalizedMessage(),e);
                        }
                        */
                  }
               }
               else if (SAN_Type.RFC822_NAME.equals(type)
                        || SAN_Type.DNS_NAME.equals(type))
               {
                  if (value instanceof String)
                  {
                     stringValue = (String) value;
                  }
               }
               else if (SAN_Type.X400_ADDRESS.equals(type))
               {
                  // TODO:
                  if (value instanceof String)
                  {
                     stringValue = (String) value;
                  }
               }
               else if (SAN_Type.DIRECTORY_NAME.equals(type))
               {
                  // TODO:
                  if (value instanceof String)
                  {
                     stringValue = (String) value;
                  }
               }
               else if (SAN_Type.EDI_PARTY_NAME.equals(type))
               {
                  // TODO:
                  if (value instanceof String)
                  {
                     stringValue = (String) value;
                  }
               }
               else if (SAN_Type.URI.equals(type))
               {
                  // TODO:
                  if (value instanceof String)
                  {
                     stringValue = (String) value;
                  }
               }
               else if (SAN_Type.IP_ADDRESS.equals(type))
               {
                  // TODO:
                  if (value instanceof String)
                  {
                     stringValue = (String) value;
                  }
               }
               else if (SAN_Type.REGISTERED_ID.equals(type))
               {
                  // TODO:
                  if (value instanceof String)
                  {
                     stringValue = (String) value;
                  }
               }

               values.add(new SAN(type, stringValue));
            }
         }
      }
      catch (CertificateParsingException e)
      {
         throw new CertificateParsingException("Error parsing SubjectAltName in certificate: " + inCertificate,e);
      }

      return values;
   }

   private static class SavingTrustManager implements X509TrustManager
   {
      private X509TrustManager mParentTrustMgr;
      private Set<X509Certificate> mAllCertificates = new HashSet<>(3);
      private Set<X509Certificate> mNewCertificates = new HashSet<>(3);


      //-----------------------------------------------------------------------
      SavingTrustManager(X509TrustManager inParentTrustMgr)
      {
         mParentTrustMgr = inParentTrustMgr;
      }

      //-----------------------------------------------------------------------
      @Override
      public void checkClientTrusted(X509Certificate[] inCertChain, String inAuthType)
            throws CertificateException
      {
         mParentTrustMgr.checkClientTrusted(inCertChain, inAuthType);
      }

      //-----------------------------------------------------------------------
      @Override
      public void checkServerTrusted(X509Certificate[] inCertChain, String inAuthType)
         throws CertificateException
      {
         CertificateException exceptionToRethrow = null;
         // check the certificate chain against the system truststore
         try
         {
            mParentTrustMgr.checkServerTrusted(inCertChain, inAuthType);
         }
         catch (CertificateException e)
         {
            // The certificate chain was found untrustworthy

            // Check if the first certificate in the chain is not known yet stored
            if (! this.mAllCertificates.contains(inCertChain[0]))
            {
               // Save the exception to be re-thrown later if not known
               exceptionToRethrow = e;
            }
         }


         // Save the full chain to both local accumulators
         for (X509Certificate cert : inCertChain)
         {
            mAllCertificates.add(cert);
            mNewCertificates.add(cert);
         }

         // check and re-throw the exception if any
         if (exceptionToRethrow != null)
         {
            throw exceptionToRethrow;
         }
      }

      //-----------------------------------------------------------------------
      @Override
      public X509Certificate[] getAcceptedIssuers()
      {
         return mParentTrustMgr.getAcceptedIssuers();
      }
   }
}