/*
 * Saml11HOKCallbackHandler.java
 *
 * Created on August 18, 2007, 11:34 PM
 */

package samlcb;

import java.io.*;
import java.util.*;
import java.math.BigInteger;

import java.security.KeyStore;
import java.security.PrivateKey;
import java.security.PublicKey;
import java.security.cert.X509Certificate;
import java.security.cert.Certificate;

import javax.security.auth.callback.Callback;
import javax.security.auth.callback.CallbackHandler;
import javax.security.auth.callback.UnsupportedCallbackException;

import com.sun.xml.wss.impl.callback.*;

import com.sun.xml.wss.saml.*;

import javax.xml.parsers.DocumentBuilderFactory;

import org.w3c.dom.*;

import com.sun.org.apache.xml.internal.security.keys.KeyInfo;

/**
 *
 * @author  Frank
 * @version
 *
 * SAML 1.1 Holder Of Key Callback Handler
 */
public  class Saml11HOKCallbackHandler implements CallbackHandler {
    
    private String keyStoreURL;
    private String keyStorePassword;
    private String keyStoreType;
    
    private String trustStoreURL;
    private String trustStorePassword;
    private String trustStoreType;
    
    private KeyStore keyStore;
    private KeyStore trustStore;
    
    private static final String fileSeparator = System.getProperty("file.separator");
    
    private  UnsupportedCallbackException unsupported =
            new UnsupportedCallbackException(null,
            "Unsupported Callback Type Encountered");
    
    private  static Element hokAssertion = null;
    
    public static final String holderOfKeyConfirmation =
            "urn:oasis:names:tc:SAML:1.0:cm:holder-of-key";
    
    String glassfish_home = System.getProperty("com.sun.aas.installRoot");
    String client_priv_key_alias="xws-security-client";
    
    public Saml11HOKCallbackHandler() {
        try {
            
            this.keyStoreURL = glassfish_home + fileSeparator + "domains" + fileSeparator + "domain1" +
                    fileSeparator + "config" + fileSeparator + "keystore.jks";
            this.keyStoreType = "JKS";
            this.keyStorePassword = "adminadmin";
            
            this.trustStoreURL = glassfish_home + fileSeparator + "domains" + fileSeparator + "domain1" +
                    fileSeparator + "config" + fileSeparator + "cacerts.jks";
            this.trustStoreType = "JKS";
            this.trustStorePassword = "adminadmin";
            
            initKeyStore();
            initTrustStore();
        }catch(Exception e) {
            e.printStackTrace();
            throw new RuntimeException(e);
        }
    }
    
    public void handle(Callback[] callbacks) throws IOException, UnsupportedCallbackException {
        for (int i=0; i < callbacks.length; i++) {
            if (callbacks[i] instanceof SAMLCallback) {
                try{
                    SAMLCallback samlCallback = (SAMLCallback)callbacks[i];
                    samlCallback.setConfirmationMethod(samlCallback.HOK_ASSERTION_TYPE);
                    if (samlCallback.getConfirmationMethod().equals(samlCallback.HOK_ASSERTION_TYPE)){
                        samlCallback.setAssertionElement(createHOKSAMLAssertion());
                        hokAssertion=samlCallback.getAssertionElement();
                    }else{
                        throw new Exception("SAML Assertion Type is not matched.");
                    }
                }catch(Exception ex){
                    ex.printStackTrace();
                }
            } else {
                throw unsupported;
            }
        }
    }
    
    private  Element createHOKSAMLAssertion() {
        
        Assertion assertion = null;
        try {
            SAMLAssertionFactory factory = SAMLAssertionFactory.newInstance(SAMLAssertionFactory.SAML1_1);
            // create the assertion id
            String assertionID = String.valueOf(System.currentTimeMillis());
            String issuer = "CN=Assertion Issuer,OU=AI,O=Assertion Issuer,L=Waltham,ST=MA,C=US";
            
            GregorianCalendar c = new GregorianCalendar();
            long beforeTime = c.getTimeInMillis();
            // roll the time by one hour
            long offsetHours = 60*60*1000;
            
            c.setTimeInMillis(beforeTime - offsetHours);
            GregorianCalendar before= (GregorianCalendar)c.clone();
            
            c = new GregorianCalendar();
            long afterTime = c.getTimeInMillis();
            c.setTimeInMillis(afterTime + offsetHours);
            GregorianCalendar after = (GregorianCalendar)c.clone();
            
            GregorianCalendar issueInstant = new GregorianCalendar();
            // statements
            List statements = new LinkedList();
            NameIdentifier nmId =
                    factory.createNameIdentifier(
                    "CN=SAML User,OU=SU,O=SAML User,L=Los Angeles,ST=CA,C=US",
                    null, "urn:oasis:names:tc:SAML:1.1:nameid-format:X509SubjectName");
            
            //default priv key cert req
            SignatureKeyCallback.DefaultPrivKeyCertRequest request =
                    new SignatureKeyCallback.DefaultPrivKeyCertRequest();
            getDefaultPrivKeyCert(request);
            
            if ( request.getX509Certificate() == null ) {
                throw new RuntimeException("Not able to resolve the Default Certificate");
            }
            PublicKey pubKey = request.getX509Certificate().getPublicKey();
            PrivateKey privKey = request.getPrivateKey();
            
            DocumentBuilderFactory docFactory = DocumentBuilderFactory.newInstance();
            
            Document doc = docFactory.newDocumentBuilder().newDocument();
            
            KeyInfo keyInfo = new KeyInfo(doc);
            keyInfo.addKeyValue(pubKey);
            
            List subConfirmation = new ArrayList();
            subConfirmation.add(holderOfKeyConfirmation);
            
            SubjectConfirmation scf =
                    factory.createSubjectConfirmation(subConfirmation, null, keyInfo.getElement());
            
            
            Subject subj = factory.createSubject(nmId, scf);
            
            List attributes = new LinkedList();
            List attributeValues = new LinkedList();
            attributeValues.add("ATTRIBUTE1");
            attributes.add( factory.createAttribute(
                    "attribute1",
                    "urn:com:sun:xml:wss:attribute",
                    attributeValues));
            
            statements.add(
                    factory.createAttributeStatement(subj, attributes));
            
            Conditions conditions = factory.createConditions(before, after, null, null, null);
            
            assertion = factory.createAssertion(assertionID, issuer, issueInstant,
                    conditions, null, statements);
            assertion.setMajorVersion(BigInteger.ONE);
            assertion.setMinorVersion(BigInteger.ONE);
            
            return assertion.sign(pubKey, privKey);
        } catch (Exception e) {
            e.printStackTrace();
            throw new RuntimeException(e);
        }
    }
    
    private void initKeyStore() throws IOException {
        try {
            keyStore = KeyStore.getInstance(keyStoreType);
            keyStore.load(new FileInputStream(keyStoreURL), keyStorePassword.toCharArray());
        } catch (Exception e) {
            throw new IOException(e.getMessage());
        }
    }
    
    private void initTrustStore() throws IOException {
        try {
            trustStore = KeyStore.getInstance(trustStoreType);
            trustStore.load(new FileInputStream(trustStoreURL), trustStorePassword.toCharArray());
        } catch (Exception e) {
            throw new IOException(e.getMessage());
        }
    }
    
    private void getDefaultPrivKeyCert(
            SignatureKeyCallback.DefaultPrivKeyCertRequest request)
            throws IOException {
        
        String uniqueAlias = null;
        try {
            Enumeration aliases = keyStore.aliases();
            while (aliases.hasMoreElements()) {
                String currentAlias = (String) aliases.nextElement();
                if (currentAlias.equals(client_priv_key_alias)){
                    if (keyStore.isKeyEntry(currentAlias)) {
                        Certificate thisCertificate = keyStore.getCertificate(currentAlias);
                        if (thisCertificate != null) {
                            if (thisCertificate instanceof X509Certificate) {
                                if (uniqueAlias == null) {
                                    uniqueAlias = currentAlias;
                                    break;
                                }
                            }
                        }
                    }
                }
            }
            if (uniqueAlias != null) {
                request.setX509Certificate(
                        (X509Certificate) keyStore.getCertificate(uniqueAlias));
                request.setPrivateKey(
                        (PrivateKey) keyStore.getKey(uniqueAlias, keyStorePassword.toCharArray()));
            }
        } catch (Exception e) {
            throw new IOException(e.getMessage());
        }
    }
}

