/*
 * JBoss, the OpenSource J2EE WebOS
 *
 * Distributable under LGPL license.
 * See terms of license at gnu.org.
 */
package org.jboss.security.plugins;

import java.io.BufferedReader;
import java.io.File;
import java.io.IOException;
import java.io.InputStream;
import java.io.InputStreamReader;
import java.lang.reflect.Method;
import java.lang.reflect.Constructor;
import java.net.MalformedURLException;
import java.net.URL;
import java.security.KeyStore;
import java.security.Security;
import java.util.Arrays;
import java.util.StringTokenizer;
import javax.crypto.Cipher;
import javax.crypto.SecretKey;
import javax.crypto.SecretKeyFactory;
import javax.crypto.spec.PBEKeySpec;
import javax.crypto.spec.PBEParameterSpec;
import javax.management.MBeanServer;
import javax.management.ObjectName;
import javax.security.auth.callback.CallbackHandler;

import com.sun.net.ssl.KeyManagerFactory;
import com.sun.net.ssl.TrustManagerFactory;
import org.jboss.mx.util.MBeanServerLocator;
import org.jboss.security.SecurityDomain;
import org.jboss.security.Util;
import org.jboss.security.auth.callback.SecurityAssociationHandler;

/** The JaasSecurityDomain is an extension of JaasSecurityManager that addes
 the notion of a KeyStore, and JSSE KeyManagerFactory and TrustManagerFactory
 for supporting SSL and other cryptographic use cases.

 Attributes:
 <ul>
   <li>KeyStoreType: The implementation type name being used, defaults to 'JKS'.
 </li>

   <li>KeyStoreURL: Set the KeyStore database URL string. This is used to obtain
  an InputStream to initialize the KeyStore. If the string is not a value
 URL, its treated as a file.
   </li>
 
<li>KeyStorePass: the password used to load the KeyStore. Its format is one of:
   <ul>
      <li>The plaintext password for the KeyStore(or whatever format is used
 by the KeyStore). The toCharArray() value of the string is used without any
 manipulation.
      </li>
      <li>A command to execute to obtain the plaintext password. The format
 is '{EXT}...' where the '...' is the exact command line that will be passed
 to the Runtime.exec(String) method to execute a platform command. The first
 line of the command output is used as the password.
      </li>
      <li>A class to create to obtain the plaintext password. The format
 is '{CLASS}classname[:ctorarg]' where the '[:ctorarg]' is an optional
 string delimited by the ':' from the classname that will be passed to the
 classname ctor. The password is obtained from classname by invoking a 'char[]
 toCharArray()' method if found, otherwise, the 'String toString()' method is
 used.
      </li> 
   </ul>
 The KeyStorePass is also used in combination with the Salt and IterationCount
 attributes to create a PBE secret key used with the encode/decode operations.
 </li>

<li>ManagerServiceName: The JMX object name string of the security manager service
 that the domain registers with to function as a security manager for the
 security domain name passed to the ctor. The makes the JaasSecurityDomain
 available under the standard JNDI java:/jaas/(domain) binding.
 </li>

<li>LoadSunJSSEProvider: A flag indicating if the Sun com.sun.net.ssl.internal.ssl.Provider 
  security provider should be loaded on startup. This is needed when using
  the Sun JSSE jars without them installed as an extension with JDK 1.3. This
  should be set to false with JDK 1.4 or when using an alternate JSSE provider
</li>

<li>Salt:
 </li>
 
<li>IterationCount:
 </li>
</ul>

 @author Scott.Stark@jboss.org
 @version $Revision: 1.6.2.11 $
*/
public class JaasSecurityDomain
   extends JaasSecurityManager
   implements SecurityDomain, JaasSecurityDomainMBean
{
   private static boolean addedSunJSSEProvider;
   /** The permission required to access encode, encode64 */
   private static final RuntimePermission encodePermission =
      new RuntimePermission("org.jboss.security.plugins.JaasSecurityDomain.encode");
   /** The permission required to access decode, decode64 */
   private static final RuntimePermission decodePermission =
      new RuntimePermission("org.jboss.security.plugins.JaasSecurityDomain.decode");

   /** The KeyStore associated with the security domain */
   private KeyStore keyStore;
   private KeyManagerFactory keyMgr;
   /** The KeyStore implementation type which defaults to 'JKS' */
   private String keyStoreType = "JKS";
   /** The resource for the keystore location */ 
   private URL keyStoreURL;
   /** The keystore password for loading */
   private char[] keyStorePassword;
   /** A command string to execute to obtain the keyStorePassword */
   private String keyStorePasswordCmd;
   /** The type of command string: EXT, CLASS */
   private String keyStorePasswordCmdType;
   /** */
   private SecretKey cipherKey;
   /** */
   private String cipherAlgorithm = "PBEwithMD5andDES";
   private byte[] salt = {1, 2, 3, 4, 5, 6, 7, 8};
   private int iterationCount = 103;
   private PBEParameterSpec cipherSpec;
   /** The JMX object name of the security manager service */
   private ObjectName managerServiceName = JaasSecurityManagerServiceMBean.OBJECT_NAME;
   /** A flag indicating if com.sun.net.ssl.internal.ssl.Provider should be installed */
   private boolean loadSunJSSEProvider = true;

   /** Creates a default JaasSecurityDomain for with a securityDomain
    name of 'other'.
    */
   public JaasSecurityDomain()
   {
      super();
   }
   /** Creates a JaasSecurityDomain for with a securityDomain
    name of that given by the 'securityDomain' argument.
    @param securityDomain , the name of the security domain
    */
   public JaasSecurityDomain(String securityDomain)
   {
      this(securityDomain, new SecurityAssociationHandler());
   }
   /** Creates a JaasSecurityDomain for with a securityDomain
    name of that given by the 'securityDomain' argument.
    @param securityDomain , the name of the security domain
    @param handler , the CallbackHandler to use to obtain login module info
    */
   public JaasSecurityDomain(String securityDomain, CallbackHandler handler)
   {
      super(securityDomain, handler);
   }

   public KeyStore getKeyStore() throws SecurityException
   {
      return keyStore;
   }
   public KeyManagerFactory getKeyManagerFactory() throws SecurityException
   {
      return keyMgr;
   }

   public KeyStore getTrustStore() throws SecurityException
   {
      return null;
   }
   public TrustManagerFactory getTrustManagerFactory() throws SecurityException
   {
      return null;
   }
   /** The JMX object name string of the security manager service.
    @return The JMX object name string of the security manager service.
    */
   public ObjectName getManagerServiceName()
   {
      return this.managerServiceName;
   }
   /** Set the JMX object name string of the security manager service.
    */
   public void setManagerServiceName(ObjectName managerServiceName)
   {
      this.managerServiceName = managerServiceName;
   }

   public String getKeyStoreType()
   {
      return this.keyStoreType;
   }
   public void setKeyStoreType(String type)
   {
      this.keyStoreType = type;
   }
   public String getKeyStoreURL()
   {
      String url = null;
      if( keyStoreURL != null )
         url = keyStoreURL.toExternalForm();
      return url;
   }
   public void setKeyStoreURL(String storeURL) throws IOException
   {
      keyStoreURL = null;
      // First see if this is a URL
      try
      {
         keyStoreURL = new URL(storeURL);
      }
      catch(MalformedURLException e)
      {
         // Not a URL or a protocol without a handler
      }

      // Next try to locate this as file path
      if( keyStoreURL == null )
      {
         File tst = new File(storeURL);
         if( tst.exists() == true )
            keyStoreURL = tst.toURL();
      }

      // Last try to locate this as a classpath resource
      if( keyStoreURL == null )
      {
         ClassLoader loader = SubjectActions.getContextClassLoader();
         keyStoreURL = loader.getResource(storeURL);
      }

      // Fail if no valid key store was located
      if( keyStoreURL == null )
      {
         String msg = "Failed to find url="+storeURL+" as a URL, file or resource";
         throw new MalformedURLException(msg);
      }
      log.debug("Using KeyStore="+keyStoreURL.toExternalForm());
   }
   public void setKeyStorePass(String password)
   {
      this.keyStorePassword = null;
      // Look for a {...} prefix indicating a password command
      if( password.charAt(0) == '{' )
      {
         StringTokenizer tokenizer = new StringTokenizer(password, "{}");
         this.keyStorePasswordCmdType = tokenizer.nextToken();
         this.keyStorePasswordCmd = tokenizer.nextToken();
      }
      else
      {
         // Its just the keystore password string
         this.keyStorePassword = password.toCharArray();
      }
   }

   public void setSalt(String salt)
   {
      this.salt = salt.getBytes();
   }

   public void setIterationCount(int iterationCount)
   {
      this.iterationCount = iterationCount;
   }

   public String getCipherAlgorithm()
   {
      return cipherAlgorithm;
   }
   public void setCipherAlgorithm(String cipherAlgorithm)
   {
      this.cipherAlgorithm = cipherAlgorithm;
   }

   public boolean getLoadSunJSSEProvider()
   {
      return loadSunJSSEProvider;
   }
   public void setLoadSunJSSEProvider(boolean flag)
   {
      this.loadSunJSSEProvider = flag;
   }

   public String getName()
   {
      return "JaasSecurityDomain("+getSecurityDomain()+")";
   }   

   /** Encrypt the secret using the cipherKey.
    * @param secret - the plaintext secret to encrypt
    * @return the encrypted secret
    * @throws Exception
    */ 
   public byte[] encode(byte[] secret)
      throws Exception
   {
      SecurityManager sm = System.getSecurityManager();
      if( sm != null )
      {
         System.out.println("Checking: "+encodePermission);
         sm.checkPermission(encodePermission);
      }

      Cipher cipher = Cipher.getInstance(cipherAlgorithm);
      cipher.init(Cipher.ENCRYPT_MODE, cipherKey, cipherSpec);
      byte[] encoding = cipher.doFinal(secret);
      return encoding;
   }
   /** Decrypt the secret using the cipherKey.
    * 
    * @param secret - the encrypted secret to decrypt.
    * @return the decrypted secret
    * @throws Exception
    */ 
   public byte[] decode(byte[] secret)
      throws Exception
   {
      SecurityManager sm = System.getSecurityManager();
      if( sm != null )
         sm.checkPermission(decodePermission);

      Cipher cipher = Cipher.getInstance(cipherAlgorithm);
      cipher.init(Cipher.DECRYPT_MODE, cipherKey, cipherSpec);
      byte[] decode = cipher.doFinal(secret);
      return decode;
   }
   /** Encrypt the secret using the cipherKey and return a base64 encoding.
    * @param secret - the plaintext secret to encrypt
    * @return the encrypted secret as a base64 string
    * @throws Exception
    */ 
   public String encode64(byte[] secret)
      throws Exception
   {
      byte[] encoding = encode(secret);
      String b64 = Util.tob64(encoding);
      return b64;
   }
   /** Decrypt the base64 encoded secret using the cipherKey.
    * 
    * @param secret - the base64 encoded encrypted secret to decrypt.
    * @return the decrypted secret
    * @throws Exception
    */
   public byte[] decode64(String secret)
      throws Exception
   {
      byte[] encoding = Util.fromb64(secret);
      byte[] decode = decode(encoding);
      return decode;
   }

   /** 
        Reload the Keystore.
   */
   public void reloadKeyStore()
      throws Exception
   {
      loadKeystore();
   }

   protected void startService() throws Exception
   {
      // Install the Sun JSSE provider unless
      synchronized( JaasSecurityDomain.class )
      {
         if( loadSunJSSEProvider == true && addedSunJSSEProvider == false )
         {
            log.debug("Adding com.sun.net.ssl.internal.ssl.Provider");
            try
            {
               addedSunJSSEProvider = Security.addProvider(new com.sun.net.ssl.internal.ssl.Provider()) != -1;
            }
            catch(Exception e)
            {
               log.warn("Failed to addProvider com.sun.net.ssl.internal.ssl.Provider", e);
            }
         }
      }

      // Load the keystore password if it was 
      loadKeystorePassword();

      // Load the keystore
      loadKeystore();
      
      /* Register with the JaasSecurityManagerServiceMBean. This allows this
       JaasSecurityDomain to function as the security manager for security-domain
       elements that declare java:/jaas/xxx for our security domain name.
       */
      MBeanServer server = MBeanServerLocator.locateJBoss();
      Object[] params = {getSecurityDomain(), this};
      String[] signature = new String[] {"java.lang.String", "org.jboss.security.SecurityDomain"};
      server.invoke(managerServiceName, "registerSecurityDomain", params, signature);
   }
   protected void stopService()
   {
      // Uninstall the Sun JSSE provider unless
      synchronized( JaasSecurityDomain.class )
      {
         if( loadSunJSSEProvider == true && addedSunJSSEProvider == true )
         {
            log.debug("Removing com.sun.net.ssl.internal.ssl.Provider");
            try
            {
               String name = (new com.sun.net.ssl.internal.ssl.Provider()).getName();
               Security.removeProvider(name);
               addedSunJSSEProvider = false;
            }
            catch(Exception e)
            {
               log.warn("Failed to removeProvider com.sun.net.ssl.internal.ssl.Provider", e);
            }
         }
      }
      
      if( keyStorePassword != null )
      {
         Arrays.fill(keyStorePassword, '\0');
         keyStorePassword = null;
      }
      cipherKey = null;
   }

   /** If keyStorePassword is null and keyStorePasswordCmd exists,
    * execute it to obtain the password.
    */ 
   private void loadKeystorePassword()
      throws Exception
   {
      if( keyStorePassword == null )
      {
         if( keyStorePasswordCmdType.equals("EXT") )
            execPasswordCmd();
         else if( keyStorePasswordCmdType.equals("CLASS") )
            invokePasswordClass();
         else
            throw new IllegalArgumentException("Unknown keyStorePasswordCmdType: "+keyStorePasswordCmdType);
      }

      // Create the PBE secret key
      cipherSpec = new PBEParameterSpec(salt, iterationCount);
      PBEKeySpec keySpec = new PBEKeySpec(keyStorePassword);
      SecretKeyFactory factory = SecretKeyFactory.getInstance("PBEwithMD5andDES");
      cipherKey = factory.generateSecret(keySpec);
   }

   private void loadKeystore()
      throws Exception
   {
      if( keyStoreURL != null )
      {
         keyStore = KeyStore.getInstance(keyStoreType);
         InputStream is = keyStoreURL.openStream();
         keyStore.load(is, keyStorePassword);
         String algorithm = KeyManagerFactory.getDefaultAlgorithm();
         keyMgr = KeyManagerFactory.getInstance(algorithm);
         keyMgr.init(keyStore, keyStorePassword);
      }
   }

   private void execPasswordCmd()
      throws Exception
   {
      log.debug("Executing command: "+keyStorePasswordCmd);
      Runtime rt = Runtime.getRuntime();
      Process p = rt.exec(keyStorePasswordCmd);
      InputStream stdin = p.getInputStream();
      BufferedReader reader = new BufferedReader(new InputStreamReader(stdin));
      String password = reader.readLine();
      stdin.close();
      int exitCode = p.waitFor();
      log.debug("Command exited with: "+exitCode);
      keyStorePassword = password.toCharArray();
   }
   /**
    * 
    * @throws Exception
    */ 
   private void invokePasswordClass()
      throws Exception
   {
      keyStorePassword = null;

      // Check for a ctor argument delimited by ':'
      String classname = keyStorePasswordCmd;
      String ctorArg = null;
      int colon = keyStorePasswordCmd.indexOf(':');
      if( colon > 0 )
      {
         classname = keyStorePasswordCmd.substring(0, colon);
         ctorArg = keyStorePasswordCmd.substring(colon+1);
      }
      log.debug("Loading class: "+classname+", ctorArg="+ctorArg);
      ClassLoader loader = SubjectActions.getContextClassLoader();
      Class c = loader.loadClass(classname);
      Object instance = null;
      // Check for a ctor(String) if ctorArg is not null
      if( ctorArg != null )
      {
         Class[] sig = {String.class};
         Constructor ctor = c.getConstructor(sig);
         Object[] args = {ctorArg};
         instance = ctor.newInstance(args);
      }
      else
      {
         // Use the default ctor
         instance = c.newInstance();
      }

      // Look for a toCharArray() method
      try
      {
         log.debug("Checking for toCharArray");
         Class[] sig = {};
         Method toCharArray = c.getMethod("toCharArray", sig);
         Object[] args = {};
         log.debug("Invoking toCharArray");
         keyStorePassword = (char[]) toCharArray.invoke(instance, args);
      }
      catch(NoSuchMethodException e)
      {
         log.debug("No toCharArray found, invoking toString");
         String tmp = instance.toString();
         if( tmp != null )
            keyStorePassword = tmp.toCharArray();
      }
   }
}
