/**
 * JOnAS: Java(TM) Open Application Server
 * Copyright (C) 2006 Bull S.A.S.
 * Contact: jonas-team@objectweb.org
 *
 * 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 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
 *
 * --------------------------------------------------------------------------
 * $Id$
 * --------------------------------------------------------------------------
 */

package org.objectweb.jonas.security.auth.spi;

import java.io.BufferedInputStream;
import java.io.File;
import java.io.FileInputStream;
import java.io.FileNotFoundException;
import java.io.IOException;
import java.security.InvalidKeyException;
import java.security.KeyStore;
import java.security.KeyStoreException;
import java.security.NoSuchAlgorithmException;
import java.security.Principal;
import java.security.PrivateKey;
import java.security.Signature;
import java.security.SignatureException;
import java.security.UnrecoverableKeyException;
import java.security.acl.Group;
import java.security.cert.CertificateException;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Enumeration;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.Set;

import javax.security.auth.Subject;
import javax.security.auth.callback.CallbackHandler;
import javax.security.auth.login.LoginException;
import javax.security.auth.spi.LoginModule;

import org.objectweb.jonas.security.auth.JSigned;


/**
 * This class is used to sign the current subject.
 * @author Florent Benoit
 */
public class SignLoginModule implements LoginModule {

    /**
     * Subject used.
     */
    private Subject subject = null;

    /**
     * Options for this login module.
     */
    private Map options = null;

    /**
     * Private key.
     */
    private static PrivateKey privateKey = null;

    /**
     * Initialize this LoginModule.
     * This method is called by the LoginContext after this LoginModule has been instantiated. The purpose of this method is to initialize this LoginModule with the relevant information. If this LoginModule does not understand any of the data stored in sharedState or options parameters, they can be ignored.
     * @param subject the Subject to be authenticated.
     * @param callbackHandler a CallbackHandler for communicating with the end user (prompting for usernames and passwords, for example).
     * @param sharedState state shared with other configured LoginModules.
     * @param options options specified in the login Configuration for this particular LoginModule.
     */
    public void initialize(Subject subject, CallbackHandler callbackHandler, Map sharedState, Map options) {
        this.subject = subject;
        this.options = options;

    }


    /**
     * Method to authenticate a Subject (phase 1).
     * The implementation of this method authenticates a Subject. For example, it may prompt for Subject information such as a username and password and then attempt to verify the password. This method saves the result of the authentication attempt as private state within the LoginModule.
     * @return true if the authentication succeeded, or false if this LoginModule should be ignored.
     * @throws LoginException if the authentication fails
     */
    public boolean login() throws LoginException {

        // Load once the private key
        if (privateKey == null) {
            // Keystore file
            String keystoreFile = (String) options.get("keystoreFile");
            if (keystoreFile == null) {
                throw new LoginException("The 'keystoreFile' attribute was not found but this attribute is mandatory");
            }

            // Keystore pass
            String keystorePass = (String) options.get("keystorePass");
            if (keystorePass == null) {
                throw new LoginException("The 'keystorePass' attribute was not found but this attribute is mandatory");
            }

            // keyPass (for the private key)
            String keyPass = (String) options.get("keyPass");
            if (keyPass == null) {
                throw new LoginException("The 'keyPass' attribute was not found but this attribute is mandatory");
            }

            // Alias
            String alias = (String) options.get("alias");
            if (alias == null) {
                throw new LoginException("The 'alias' attribute was not found but this attribute is mandatory");
            }


            // Check that the file exists
            File f = new File(keystoreFile);
            if (!f.exists()) {
                throw new LoginException("The keystore file named '" + f + "' was not found.");
            }

            // Gets the keystore instance
            KeyStore keyStore = null;
            try {
                keyStore = KeyStore.getInstance(KeyStore.getDefaultType());
            } catch (KeyStoreException e) {
                throw new LoginException("Error while getting a keystore '" + e.getMessage());
            }

            // Load the keystore file
            try {
                keyStore.load(new BufferedInputStream(new FileInputStream(f)), keystorePass.toCharArray());
            } catch (NoSuchAlgorithmException e) {
                throw new LoginException("Error while loading the keystore file '" + f + "'." + e.getMessage());
            } catch (CertificateException e) {
                throw new LoginException("Error while loading the keystore file '" + f + "'." + e.getMessage());
            } catch (FileNotFoundException e) {
                throw new LoginException("Error while loading the keystore file '" + f + "'." + e.getMessage());
            } catch (IOException e) {
                throw new LoginException("Error while loading the keystore file '" + f + "'." + e.getMessage());
            }

            // Get the private key
            try {
                privateKey = (PrivateKey) keyStore.getKey(alias, keyPass.toCharArray());
            } catch (KeyStoreException e) {
                throw new LoginException("Error while getting alias named '" + alias + "' in the keystore file '" + f + "'." + e.getMessage());
            } catch (NoSuchAlgorithmException e) {
                throw new LoginException("Error while getting alias named '" + alias + "' in the keystore file '" + f + "'." + e.getMessage());
            } catch (UnrecoverableKeyException e) {
                throw new LoginException("Error while getting alias named '" + alias + "' in the keystore file '" + f + "'." + e.getMessage());
            }
        }

        return true;
    }


    /**
     * Method to commit the authentication process (phase 2).
     * This method is called if the LoginContext's overall authentication succeeded (the relevant REQUIRED, REQUISITE, SUFFICIENT and OPTIONAL LoginModules succeeded).
     * If this LoginModule's own authentication attempt succeeded (checked by retrieving the private state saved by the login method), then this method associates relevant Principals and Credentials with the Subject located in the LoginModule. If this LoginModule's own authentication attempted failed, then this method removes/destroys any state that was originally saved.
     * @return true if this method succeeded, or false if this LoginModule should be ignored.
     * @throws LoginException if the commit fails
     */
    public boolean commit() throws LoginException {


        // Get Principal Name
        String principalName = null;

        // Retrieve only principal name (without groups)
        Set principals = subject.getPrincipals(Principal.class);
        Iterator iterator = principals.iterator();
        while (iterator.hasNext()) {
            Principal principal = (Principal) iterator.next();
            if (!(principal instanceof Group)) {
               principalName = principal.getName();
            }
        }

        // No name --> error
        if (principalName == null) {
            throw new LoginException("There was no previous login module. This login module can only be used in addition to another module which perform the authentication.");
        }

        // Retrieve all roles of the user (Roles are members of the Group.class)
        List principalRoles = new ArrayList();
        principals = subject.getPrincipals(Group.class);
        iterator = principals.iterator();
        while (iterator.hasNext()) {
            Group group = (Group) iterator.next();
            Enumeration e = group.members();
            while (e.hasMoreElements()) {
                Principal p = (Principal) e.nextElement();
                principalRoles.add(p.getName());
            }
        }


        // Build signature
        Signature signature = null;
        try {
            signature = Signature.getInstance("SHA1withDSA");
        } catch (NoSuchAlgorithmException e) {
           throw new LoginException("The algorithm 'SHA1withDSA' was not found:" + e.getMessage());
        }
        try {
            signature.initSign(privateKey);
        } catch (InvalidKeyException e) {
            throw new LoginException("Invalid private key when initializing signature:" + e.getMessage());
        }

        // Add principal name
        try {
            signature.update(principalName.getBytes());
        } catch (SignatureException e) {
            throw new LoginException("Cannot add the bytes of the principal name:" + e.getMessage());
        }

        // Get roles
        String[] roles = (String[]) principalRoles.toArray(new String[principalRoles.size()]);
        // Sort roles before adding it to the signature to preserve the order
        Arrays.sort(roles);

        // Add each role
        for (int i = 0; i < roles.length; i++) {
            try {
                signature.update(roles[i].getBytes());
            } catch (SignatureException e) {
                throw new LoginException("Cannot add the bytes of one role:" + e.getMessage());
            }
        }

        // Get bytes of this signature
        byte[] signedData = null;
        try {
            signedData = signature.sign();
        } catch (SignatureException e) {
            throw new LoginException("Cannot sign the data:" + e.getMessage());
        }

        // Store the signature
        JSigned signed = new JSigned(signedData);
        subject.getPrincipals().add(signed);

        return true;
    }


    /**
     * Method to abort the authentication process (phase 2).
     * This method is called if the LoginContext's overall authentication failed. (the relevant REQUIRED, REQUISITE, SUFFICIENT and OPTIONAL LoginModules did not succeed).
     * If this LoginModule's own authentication attempt succeeded (checked by retrieving the private state saved by the login method), then this method cleans up any state that was originally saved.
     * @return true if this method succeeded, or false if this LoginModule should be ignored.
     * @throws LoginException if the abort fails
     */
    public boolean abort() throws LoginException {
        // Do nothing (as all is done in the commit() phase)
        return true;
    }

    /**
     * Method which logs out a Subject.
     * An implementation of this method might remove/destroy a Subject's Principals and Credentials.
     * @return true if this method succeeded, or false if this LoginModule should be ignored.
     * @throws LoginException if the logout fails
     */
    public boolean logout() throws LoginException {
        // Nothing
        return true;

    }

}
