/* 
    * Copyright (c) 2007, Fraunhofer-Gesellschaft
    * All rights reserved.
    * 
    * Redistribution and use in source and binary forms, with or without
    * modification, are permitted provided that the following conditions are
    * met:
    * 
    * (1) Redistributions of source code must retain the above copyright
   *     notice, this list of conditions and the disclaimer at the end.
   *     Redistributions in binary form must reproduce the above copyright
   *     notice, this list of conditions and the following disclaimer in
   *     the documentation and/or other materials provided with the
   *     distribution.
   * 
   * (2) Neither the name of Fraunhofer nor the names of its
   *     contributors may be used to endorse or promote products derived
   *     from this software without specific prior written permission.
   * 
   * DISCLAIMER
   * 
   * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
   * "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
   * LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
   * A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
   * OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
   * SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
   * LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
   * DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
   * THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
   * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
   * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
   *  
   */
  package org.ogf.graap.wsag.client.wsrf.security;
  
  import java.io.IOException;
  import java.io.InputStream;
  import java.security.KeyStore;
  import java.security.KeyStoreException;
  import java.security.NoSuchAlgorithmException;
  import java.security.PrivateKey;
  import java.security.cert.Certificate;
  import java.security.cert.CertificateException;
  import java.security.cert.X509Certificate;
  import java.text.MessageFormat;
  import java.util.Map;
  import java.util.Properties;
  
  import javax.security.auth.Subject;
  import javax.security.auth.callback.Callback;
  import javax.security.auth.callback.CallbackHandler;
  import javax.security.auth.callback.UnsupportedCallbackException;
  import javax.security.auth.login.LoginException;
  import javax.security.auth.spi.LoginModule;
  import javax.security.auth.x500.X500Principal;
  import javax.security.auth.x500.X500PrivateCredential;
  
  import org.apache.log4j.Logger;
  import org.apache.ws.security.WSSecurityException;
  import org.apache.ws.security.components.crypto.Crypto;
  import org.apache.ws.security.components.crypto.CryptoBase;
  import org.ogf.graap.wsag.api.configuration.WSAG4JConfiguration;
  import org.ogf.graap.wsag.api.security.KeystoreCallback;
  import org.ogf.graap.wsag.api.security.SecurityConstants;
  
  /**
   * KeystoreLoginModule
   * 
   * @author Oliver Waeldrich
   * 
   */
  public class KeystoreLoginModule
      implements LoginModule
  {
  
      private static final Logger LOG = Logger.getLogger( KeystoreLoginModule.class );
  
      private Subject klmSubject;
  
      private CallbackHandler cbHandler;
  
      // private Map sharedState;
  
      @SuppressWarnings( "rawtypes" )
      private Map klmOptions;
  
      // variables related to access the keystore
      private KeyStore keystore;
  
      private String keystoreType;
  
      private String keystoreFile;
  
      private String keystorePassword;
  
      private String alias;
  
      private String privateKeyPassword;
 
     private String truststoreType;
 
     private String truststoreFile;
 
     private String truststorePassword;
 
     // variables related to user authentication
     private Crypto userCrypto;
 
     private X500Principal userPrincipal;
 
     // variables related to login module steering
     private boolean login = false;
 
     private boolean commit = false;
 
     /**
      * {@inheritDoc}
      * 
      * @see javax.security.auth.spi.LoginModule#initialize(javax.security.auth.Subject,
      *      javax.security.auth.callback.CallbackHandler, java.util.Map, java.util.Map)
      */
     @Override
     @SuppressWarnings( "rawtypes" )
     public void initialize( Subject subject, CallbackHandler callbackHandler, Map sharedState, Map options )
     {
         this.klmSubject = subject;
         this.cbHandler = callbackHandler;
         // this.sharedState = sharedState;
         this.klmOptions = options;
 
         initializeOptions();
     }
 
     private void initializeOptions()
     {
 
         keystoreFile = (String) klmOptions.get( "keyStoreURL" );
         keystoreType = (String) klmOptions.get( "keyStoreType" );
         alias = (String) klmOptions.get( "keyStoreAlias" );
 
         truststoreFile = (String) klmOptions.get( "trustStoreURL" );
         truststoreType = (String) klmOptions.get( "trustStoreType" );
 
         //
         // processing of the options
         //
         // keystoreFile = resolveKeystoreURL(keystoreFile);
         // truststoreFile = resolveKeystoreURL(truststoreFile);
 
         keystoreType = ( keystoreType == null ) ? "JKS" : keystoreType;
         truststoreType = ( truststoreType == null ) ? "JKS" : truststoreType;
     }
 
     /**
      * {@inheritDoc}
      * 
      * @see javax.security.auth.spi.LoginModule#login()
      */
     @Override
     public boolean login() throws LoginException
     {
         KeystoreCallback ksCallback = new KeystoreCallback();
         Callback[] callbacks = new Callback[] { ksCallback };
 
         // handle login callbacks
         try
         {
             cbHandler.handle( callbacks );
         }
         catch ( IOException e )
         {
             String message = "IO error during login";
             LoginException le = new LoginException( message );
             le.initCause( e );
             throw le;
         }
         catch ( UnsupportedCallbackException e )
         {
             String message = "Invalid callback handler. Callback not supported.";
             LoginException le = new LoginException( message );
             le.initCause( e );
             throw le;
         }
 
         keystorePassword = ksCallback.getKeystorePassword();
 
         truststorePassword = ksCallback.getTruststorePassword();
 
         privateKeyPassword = ksCallback.getPrivateKeyPassword();
 
         //
         // if an empty alias is supplied, we set the alias to null
         // this is for treating PKCS12 files, where certificates
         // do not have a alias
         //
         // if ("".equals(alias)) alias = null;
 
         if ( ( keystoreFile == null ) || ( keystorePassword == null ) || ( privateKeyPassword == null ) )
         {
 
             String message =
                 "Missing required parameter. "
                     + "The KeystoreLoginModule requires the following parameters: "
                     + "[keystoreFilename, keystorePassword, alias, privateKeyPassword]";
 
             throw new LoginException( message );
         }
 
         loadKeyStore();
 
         login = true;
 
         return true;
     }
 
     /**
      * {@inheritDoc}
      * 
      * @see javax.security.auth.spi.LoginModule#commit()
      */
     @Override
     public boolean commit() throws LoginException
     {
         if ( !login )
         {
             return false;
         }
 
         PrivateKey userKey;
         X500PrivateCredential userCredential;
         X509Certificate[] userCertificateChain;
 
         try
         {
             userCertificateChain = getCertificates( alias );
             userKey = (PrivateKey) keystore.getKey( alias, privateKeyPassword.toCharArray() );
         }
         catch ( KeyStoreException e )
         {
             // thrown by keystoreManager.getCertificateByAlias(defaultAlias)[0]
             String message = "Could not get default certificate from KeyStoreManager";
             LoginException le = new LoginException( message );
             le.initCause( e );
             throw le;
         }
         catch ( Exception e )
         {
             // thrown by keystoreManager.getKeyEntry(defaultAlias)
             String message = "Could not get private key from KeyStoreManager";
             LoginException le = new LoginException( message );
             le.initCause( e );
             throw le;
         }
 
         if ( userCertificateChain == null )
         {
             Object[] filler = new Object[] { alias };
             String message = MessageFormat.format( "No certificates found for user {0}", filler );
             throw new LoginException( message );
         }
 
         userCredential = new X500PrivateCredential( userCertificateChain[0], userKey );
 
         userCrypto = loadUserCrypto( userCredential, userCertificateChain, keystore );
         userPrincipal =
             new X500Principal( userCredential.getCertificate().getSubjectX500Principal().getName() );
 
         klmSubject.getPrivateCredentials().add( userCrypto );
         klmSubject.getPrivateCredentials().add( userCredential );
         klmSubject.getPrincipals().add( userPrincipal );
 
         commit = true;
 
         return true;
     }
 
     /**
      * {@inheritDoc}
      * 
      * @see javax.security.auth.spi.LoginModule#abort()
      */
     @Override
     public boolean abort() throws LoginException
     {
         if ( !login )
         {
             return false;
         }
         if ( ( login ) && ( !commit ) )
         {
             // login succeeded, but overall authentication failed
             login = false;
 
             klmSubject.getPrincipals().remove( userPrincipal );
             klmSubject.getPrivateCredentials().remove( userCrypto );
 
             userCrypto = null;
             userPrincipal = null;
 
             keystore = null;
             keystoreFile = null;
             keystorePassword = null;
             keystoreType = null;
 
             alias = null;
             privateKeyPassword = null;
         }
         else
         {
             // overall authentication succeeded and commit succeeded,
             // but someone else's commit failed
             logout();
         }
 
         return true;
     }
 
     /**
      * {@inheritDoc}
      * 
      * @see javax.security.auth.spi.LoginModule#logout()
      */
     @Override
     public boolean logout() throws LoginException
     {
         klmSubject.getPrivateCredentials().remove( userCrypto );
         klmSubject.getPrincipals().remove( userPrincipal );
 
         login = false;
         commit = false;
 
         userCrypto = null;
         userPrincipal = null;
 
         keystore = null;
         keystoreFile = null;
         keystorePassword = null;
         keystoreType = null;
 
         alias = null;
         privateKeyPassword = null;
 
         return true;
     }
 
     private synchronized KeyStore getKeystore() throws LoginException
     {
         if ( keystore == null )
         {
             loadKeyStore();
         }
         return keystore;
     }
 
     private void loadKeyStore() throws LoginException
     {
         try
         {
             String actualKSType = ( keystoreType == null ) ? KeyStore.getDefaultType() : keystoreType;
 
             keystore = KeyStore.getInstance( actualKSType );
 
             if ( keystoreFile == null )
             {
                 throw new IOException( "No keystore specified by user." );
             }
 
             InputStream ksInput = WSAG4JConfiguration.findResource( keystoreFile );
             keystore.load( ksInput, keystorePassword.toCharArray() );
 
         }
         catch ( KeyStoreException e )
         {
             throw new LoginException( e.getMessage() );
         }
         catch ( IOException e )
         {
             throw new LoginException( e.getMessage() );
         }
         catch ( CertificateException e )
         {
             throw new LoginException( e.getMessage() );
         }
         catch ( NoSuchAlgorithmException e )
         {
             throw new LoginException( e.getMessage() );
         }
 
     }
 
     /**
      * Gets the list of certificates for a given alias.
      * <p/>
      * 
      * @param ksAlias
      *            Lookup certificate chain for this alias
      * 
      * @return Array of X509 certificates for this alias name, or null if this alias does not exist in the
      *         keystore
      * 
      * @throws KeyStoreException
      *             error accessing the keystore
      * 
      * @throws LoginException
      *             error logging into the keystore
      */
     private X509Certificate[] getCertificates( String ksAlias ) throws KeyStoreException, LoginException
     {
         Certificate[] certs = null;
         Certificate cert = null;
 
         KeyStore store = getKeystore();
 
         if ( store != null )
         {
             // There's a chance that there can only be a set of trust stores
             certs = store.getCertificateChain( ksAlias );
             if ( certs == null || certs.length == 0 )
             {
                 // no cert chain, so lets check if getCertificate gives us a
                 // result.
                 cert = store.getCertificate( ksAlias );
             }
         }
 
         if ( cert != null )
         {
             certs = new Certificate[] { cert };
         }
         else if ( certs == null )
         {
             // At this pont we don't have certs or a cert
             return null;
         }
 
         X509Certificate[] x509certs = new X509Certificate[certs.length];
         for ( int i = 0; i < certs.length; i++ )
         {
             x509certs[i] = (X509Certificate) certs[i];
         }
         return x509certs;
     }
 
     /**
      * This method returns a new instance of org.apache.ws.security.components.crypto.Merlin.
      * 
      * @param privCredential
      * @param pubCredential
      * @param store
      * @return
      */
     private Crypto loadUserCrypto( final X500PrivateCredential privCredential,
                                    final X509Certificate[] pubCredential, final KeyStore store )
     {
 
         //
         // create the Merlin properties
         //
         Properties properties = new Properties();
         properties.setProperty( SecurityConstants.PROP_CRYPTO_PROVIDER,
             org.apache.ws.security.components.crypto.Merlin.class.getName() );
 
         properties.setProperty( SecurityConstants.PROP_KEYSTORE_TYPE, keystoreType );
         properties.setProperty( SecurityConstants.PROP_KEYSTORE_PASS, keystorePassword );
         properties.setProperty( SecurityConstants.PROP_KEYSTORE_ALIAS, alias );
         properties.setProperty( SecurityConstants.PROP_KEYSTORE_ALIAS_PASS, privateKeyPassword );
 
         if ( keystoreFile.startsWith( "/" ) )
         {
             properties.setProperty( SecurityConstants.PROP_KEYSTORE_FILE, keystoreFile.substring( 1 ) );
         }
         else
         {
             properties.setProperty( SecurityConstants.PROP_KEYSTORE_FILE, keystoreFile );
         }
 
         if ( truststoreFile != null )
         {
             properties.setProperty( SecurityConstants.PROP_TRUSTSTORE_FILE, truststoreFile );
         }
         if ( truststorePassword != null )
         {
             properties.setProperty( SecurityConstants.PROP_TRUSTSTORE_PASS, truststorePassword );
         }
         if ( truststoreType != null )
         {
             properties.setProperty( SecurityConstants.PROP_TRUSTSTORE_TYPE, truststoreType );
         }
 
         try
         {
             ClassLoader classLoader = getClass().getClassLoader();
             return new org.apache.ws.security.components.crypto.Merlin( properties, classLoader )
             {
                 /**
                  * Change the alias name provided in the AXIS2 service configuration to the defaultAlias if
                  * necessary and change the password accordingly.
                  * 
                  * {@inheritDoc}
                  */
                 @Override
                 public PrivateKey getPrivateKey( String myAlias, String myPassword ) throws Exception
                 {
                     if ( myAlias.equals( SecurityConstants.DEFAULT_ALIAS )
                         && myPassword.equals( SecurityConstants.DEFAULT_ALIAS_PASSWORD ) )
                     {
                         myAlias = getDefaultX509Alias();
                         myPassword =
                             properties.getProperty( "org.apache.ws.security.crypto.merlin.alias.password" );
                     }
                     return super.getPrivateKey( myAlias, myPassword );
                 }
 
                 /**
                  * Change the alias name provided in the AXIS2 service configuration to the defaultAlias if
                  * necessary.
                  * 
                  * {@inheritDoc}
                  */
                 @Override
                 public X509Certificate[] getCertificates( String myAlias ) throws WSSecurityException
                 {
                     myAlias =
                         ( myAlias.equals( SecurityConstants.DEFAULT_ALIAS ) ) ? getDefaultX509Alias()
                                         : myAlias;
                     return super.getCertificates( myAlias );
                 }
             };
         }
         catch ( Exception e )
         {
             Object[] filler = new Object[] { e.getMessage() };
             String message = MessageFormat.format( "Could not load user crypto. Reason: {0}", filler );
             LOG.error( message );
             LOG.error( "Try fallback... Does eventually not work for PKCS12 keystore." );
         }
 
         //
         // Fallback in case Merlin was not loaded
         //
         CryptoBase crypto = getFallbackCrypto( privCredential, pubCredential, store );
 
         return crypto;
     }
 
     /**
      * @param privCredential
      * @param pubCredential
      * @param store
      * @return
      */
     private CryptoBase getFallbackCrypto( final X500PrivateCredential privCredential,
                                           final X509Certificate[] pubCredential, final KeyStore store )
     {
         CryptoBase crypto = new CryptoBase()
         {
 
             @Override
             protected String getCryptoProvider()
             {
                 return null;
             }
 
             @Override
             public String getDefaultX509Alias()
             {
                 /*
                  * The user name for signing a SOAP message is configured statically in the AXIS2
                  * configuration of the Rampart module. By default we set it to
                  * SecurityConstants.DEFAULT_ALIAS.
                  */
                 return SecurityConstants.DEFAULT_ALIAS;
             }
 
             /**
              * {@inheritDoc}
              */
             @Override
             public PrivateKey getPrivateKey( String ksAlias, String password ) throws Exception
             {
                 /*
                  * We return the private credential for requests with the SecurityConstants.DEFAULT_ALIAS,
                  * otherwise we delegate the request to the keystore object.
                  */
                 if ( SecurityConstants.DEFAULT_ALIAS.equals( ksAlias ) )
                 {
                     return privCredential.getPrivateKey();
                 }
                 return (PrivateKey) keystore.getKey( ksAlias, password.toCharArray() );
             }
 
             /**
              * {@inheritDoc}
              */
             @Override
             public X509Certificate[] getCertificates( String ksAlias ) throws WSSecurityException
             {
                 /*
                  * We return the public credential for requests with the SecurityConstants.DEFAULT_ALIAS,
                  * otherwise we delegate the request to the keystore object.
                  */
                 if ( SecurityConstants.DEFAULT_ALIAS.equals( ksAlias ) )
                 {
                     if ( pubCredential == null )
                     {
                         LOG.warn( "No certificate chain not provided in the login context." );
                         return new X509Certificate[] { privCredential.getCertificate() };
                     }
 
                     return pubCredential;
                 }
                 else
                 {
                     return super.getCertificates( ksAlias );
                 }
             }
 
         };
 
         crypto.setKeyStore( store );
         return crypto;
     }
 
 }