/*
 *  Copyright 1999-2004 The Apache Software Foundation.
 *
 *  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.
 */

package org.jboss.axis.components.net;

import java.io.File;
import java.io.FileInputStream;
import java.io.FileNotFoundException;
import java.io.IOException;
import java.io.InputStream;
import java.net.Socket;
import java.net.URL;
import java.security.KeyStore;
import java.security.Principal;
import java.security.PrivateKey;
import java.security.SecureRandom;
import java.security.cert.X509Certificate;
import java.util.HashMap;

import javax.net.ssl.KeyManager;
import javax.net.ssl.KeyManagerFactory;
import javax.net.ssl.SSLContext;
import javax.net.ssl.TrustManager;
import javax.net.ssl.TrustManagerFactory;
import javax.net.ssl.X509KeyManager;

import org.jboss.logging.Logger;

/**
 * SSL server socket factory. It _requires_ a valid RSA key and JSSE. This is
 * derived from the tomcat source.
 *
 * @author Harish Prabandham
 * @author Costin Manolache
 * @author Stefan Freyr Stefansson
 * @author EKR -- renamed to JSSESocketFactory
 * @author Jan Luehe
 * @author <a href="jason.greene@jboss.com">Jason T. Greene</a>
 */
public class JSSE14SocketFactory extends JSSESocketFactory implements
      SecureSocketFactory
{

   private static Logger log = Logger.getLogger(JSSE14SocketFactory.class);

   /**
    * Field defaultkeyStoreType
    */
   static String defaultKeyStoreType = "JKS";

   /**
    * Field defaultProtocol
    */
   static String defaultProtocol = "TLS";

   /**
    * Field defaultAlgorithm
    */
   static String defaultAlgorithm = "SunX509";

   /**
    * Field defaultKeyPass
    */
   static String defaultKeyStorePassword = "changeit";

   /**
    * Constructor JSSESocketFactory
    *
    * @param options
    */
   public JSSE14SocketFactory(HashMap options)
   {
      super((options == null) ? new HashMap() : options);
   }

   protected String getKeyStoreType()
   {
      String keyStoreType = (String)options.get("org.jboss.webservice.keyStoreType");
      if (keyStoreType == null)
         keyStoreType = System.getProperty("javax.net.ssl.keyStoreType");

      if (keyStoreType == null)
         keyStoreType = defaultKeyStoreType;

      return keyStoreType;
   }

   protected String getTrustStorePassword()
   {
      String trustStorePassword = (String)options.get("org.jboss.webservice.trustStorePassword");
      if (trustStorePassword == null)
      {
         trustStorePassword = System
               .getProperty("javax.net.ssl.trustStorePassword");
      }
      if (trustStorePassword == null)
      {
         trustStorePassword = getKeyStorePassword();
      }
      if (log.isDebugEnabled())
      {
         log.debug("TrustPass = " + trustStorePassword);
      }
      return trustStorePassword;
   }

   protected String getTrustStoreType()
   {
      String truststoreType = (String)options.get("org.jboss.webservice.trustStoreType");
      if (truststoreType == null)
         truststoreType = System.getProperty("javax.net.ssl.trustStoreType");

      if (truststoreType == null)
         truststoreType = defaultKeyStoreType;

      if (log.isDebugEnabled())
      {
         log.debug("trustType = " + truststoreType);
      }

      return truststoreType;
   }

   /**
    * Reads the keyStore and initializes the SSL socket factory.
    */
   protected void initFactory() throws IOException
   {
      try
      {
         // SSL protocol variant (e.g., TLS, SSL v3, etc.)
         String protocol = (String)options.get("org.jboss.webservice.protocol");
         if (protocol == null)
         {
            protocol = defaultProtocol;
         }

         // Certificate encoding algorithm (e.g., SunX509)
         String algorithm = (String)options.get("org.jboss.webservice.algorithm");
         if (algorithm == null)
         {
            algorithm = defaultAlgorithm;
         }

         String trustAlgorithm = (String)options.get("org.jboss.webservice.truststoreAlgorithm");
         if (trustAlgorithm == null)
         {
            trustAlgorithm = algorithm;
         }

         String keyStoreType = getKeyStoreType();
         String trustStoreType = getTrustStoreType();

         // Create and init SSLContext
         SSLContext context = SSLContext.getInstance(protocol);
         context.init(getKeyManagers(keyStoreType, algorithm, (String)options
               .get("org.jboss.webservice.keyAlias")), getTrustManagers(trustAlgorithm,
               trustStoreType), new SecureRandom());

         // create proxy
         sslFactory = context.getSocketFactory();
      }
      catch (IOException e)
      {
         throw e;
      }
      catch (RuntimeException e)
      {
         throw e;
      }
      catch (Exception e)
      {
         throw new IOException(e.getMessage());
      }
   }

   /**
    * Gets the initialized key managers.
    */
   protected KeyManager[] getKeyManagers(String type, String algorithm,
         String keyAlias) throws Exception
   {

      KeyManager[] kms = null;

      String keyStorePass = getKeyStorePassword();

      KeyStore ks = getKeyStore(type, keyStorePass);
      if (ks == null)
         return null;
      
      if (keyAlias != null && !ks.isKeyEntry(keyAlias))
      {
         throw new IOException("Could not find alias in keyStore: " + keyAlias);
      }

      KeyManagerFactory kmf = KeyManagerFactory.getInstance(algorithm);
      kmf.init(ks, keyStorePass.toCharArray());

      kms = kmf.getKeyManagers();
      if (keyAlias != null)
      {
         if (defaultKeyStoreType.equals(type))
         {
            keyAlias = keyAlias.toLowerCase();
         }
         for (int i = 0; i < kms.length; i++)
         {
            kms[i] = new JSSEKeyManager((X509KeyManager)kms[i], keyAlias);
         }
      }

      return kms;
   }

   /*
    * Gets the SSL keystore password.
    */
   protected String getKeyStorePassword()
   {
      String keyStorePassword = (String)options.get("org.jboss.webservice.keyStorePassword");
      if (keyStorePassword == null)
         keyStorePassword = System.getProperty("javax.net.ssl.keyStorePassword");

      if (keyStorePassword == null)
         keyStorePassword = defaultKeyStorePassword;

      return keyStorePassword;
   }

   /*
    * Gets the SSL server's keystore.
    */
   protected KeyStore getKeyStore(String type, String pass) throws IOException
   {
      String keyStoreFile = (String)options.get("org.jboss.webservice.keyStore");
      if (keyStoreFile == null)
         keyStoreFile = System.getProperty("javax.net.ssl.keyStore");

      if (keyStoreFile == null)
         return null;

      return getStore(type, keyStoreFile, pass);
   }

   /**
    * Gets the intialized trust managers.
    */
   protected TrustManager[] getTrustManagers(String algorithm, String type)
         throws Exception
   {

      TrustManager[] tms = null;

      KeyStore trustStore = getTrustStore(type, getTrustStorePassword());
      if (trustStore != null)
      {
         TrustManagerFactory tmf = TrustManagerFactory.getInstance(algorithm);
         tmf.init(trustStore);
         tms = tmf.getTrustManagers();
      }

      return tms;
   }

   protected KeyStore getTrustStore(String type, String pass)
         throws IOException
   {
      KeyStore trustStore = null;

      String trustStoreFile = (String)options.get("org.jboss.webservice.trustStore");
      if (trustStoreFile == null)
      {
         trustStoreFile = System.getProperty("javax.net.ssl.trustStore");
      }
      if (log.isDebugEnabled())
      {
         log.debug("Truststore = " + trustStoreFile);
      }

      if (trustStoreFile != null && pass != null)
      {
         trustStore = getStore(type, trustStoreFile, pass);
      }

      return trustStore;
   }

   /*
    * Gets the key- or truststore with the specified type, path, and password.
    */
   private KeyStore getStore(String type, String path, String pass)
         throws IOException
   {
      KeyStore ks = null;
      InputStream istream = null;

      try
      {
         ks = KeyStore.getInstance(type);

         // Allow for keystores to be in a classloader, search them first
         URL resource = Thread.currentThread().getContextClassLoader().getResource(path);
         if (resource != null)
         {
            istream = resource.openStream();
         }
         else
         {
            File keyStoreFile = new File(path);
            istream = new FileInputStream(keyStoreFile);
         }

         ks.load(istream, pass.toCharArray());
         istream.close();
         istream = null;
      }
      catch (FileNotFoundException fnfe)
      {
         throw fnfe;
      }
      catch (IOException ioe)
      {
         throw ioe;
      }
      catch (Exception ex)
      {
         ex.printStackTrace();
         throw new IOException("Exception trying to load keystore " + path
               + ": " + ex.getMessage());
      }
      finally
      {
         if (istream != null)
         {
            try
            {
               istream.close();
            }
            catch (IOException ioe)
            {
               // Do nothing
            }
         }
      }

      return ks;
   }

   public static final class JSSEKeyManager implements X509KeyManager
   {

      private X509KeyManager delegate;

      private String clientKeyAlias;

      /**
       * Constructor.
       *
       * @param mgr
       *           The X509KeyManager used as a delegate
       * @param serverKeyAlias
       *           The alias name of the server's keypair and supporting
       *           certificate chain
       */
      public JSSEKeyManager(X509KeyManager mgr, String clientKeyAlias)
      {
         this.delegate = mgr;
         this.clientKeyAlias = clientKeyAlias;
      }

      /**
       * Choose an alias to authenticate the client side of a secure socket,
       * given the public key type and the list of certificate issuer
       * authorities recognized by the peer (if any).
       *
       * @param keyType
       *           The key algorithm type name(s), ordered with the
       *           most-preferred key type first
       * @param issuers
       *           The list of acceptable CA issuer subject names, or null if it
       *           does not matter which issuers are used
       * @param socket
       *           The socket to be used for this connection. This parameter can
       *           be null, in which case this method will return the most
       *           generic alias to use
       *
       * @return The alias name for the desired key, or null if there are no
       *         matches
       */
      public String chooseClientAlias(String[] keyType, Principal[] issuers,
            Socket socket)
      {
         return clientKeyAlias;
      }

      /**
       * Returns this key manager's server key alias that was provided in the
       * constructor.
       *
       * @param keyType
       *           The key algorithm type name (ignored)
       * @param issuers
       *           The list of acceptable CA issuer subject names, or null if it
       *           does not matter which issuers are used (ignored)
       * @param socket
       *           The socket to be used for this connection. This parameter can
       *           be null, in which case this method will return the most
       *           generic alias to use (ignored)
       *
       * @return Alias name for the desired key
       */
      public String chooseServerAlias(String keyType, Principal[] issuers,
            Socket socket)
      {
         return delegate.chooseServerAlias(keyType, issuers, socket);
      }

      /**
       * Returns the certificate chain associated with the given alias.
       *
       * @param alias
       *           The alias name
       *
       * @return Certificate chain (ordered with the user's certificate first
       *         and the root certificate authority last), or null if the alias
       *         can't be found
       */
      public X509Certificate[] getCertificateChain(String alias)
      {
         return delegate.getCertificateChain(alias);
      }

      /**
       * Get the matching aliases for authenticating the client side of a secure
       * socket, given the public key type and the list of certificate issuer
       * authorities recognized by the peer (if any).
       *
       * @param keyType
       *           The key algorithm type name
       * @param issuers
       *           The list of acceptable CA issuer subject names, or null if it
       *           does not matter which issuers are used
       *
       * @return Array of the matching alias names, or null if there were no
       *         matches
       */
      public String[] getClientAliases(String keyType, Principal[] issuers)
      {
         return delegate.getClientAliases(keyType, issuers);
      }

      /**
       * Get the matching aliases for authenticating the server side of a secure
       * socket, given the public key type and the list of certificate issuer
       * authorities recognized by the peer (if any).
       *
       * @param keyType
       *           The key algorithm type name
       * @param issuers
       *           The list of acceptable CA issuer subject names, or null if it
       *           does not matter which issuers are used
       *
       * @return Array of the matching alias names, or null if there were no
       *         matches
       */
      public String[] getServerAliases(String keyType, Principal[] issuers)
      {
         return delegate.getServerAliases(keyType, issuers);
      }

      /**
       * Returns the key associated with the given alias.
       *
       * @param alias
       *           The alias name
       *
       * @return The requested key, or null if the alias can't be found
       */
      public PrivateKey getPrivateKey(String alias)
      {
         return delegate.getPrivateKey(alias);
      }
   }

}
