Accessing Claim Aware Services using STS Secured with Non-repudiation

public class Client {
 
    ...
     
    public static void main(String[] args) {
        Client client = new Client();
        client.run();
    }
 
    private void run() {
        try {
            loadConfigurations();
 
            // set the trust store as a system property for communication over
            // TLS.
            System.setProperty("javax.net.ssl.trustStore", keystorePath);
            System.setProperty("javax.net.ssl.trustStorePassword", keystorePwd);
 
            // create configuration context
            ConfigurationContext configCtx = ConfigurationContextFactory
                    .createConfigurationContextFromFileSystem(repoPath);
 
            // create STS client
            STSClient stsClient = new STSClient(configCtx);
            stsClient.setRstTemplate(getRSTTemplate());
 
            String action = null;
            String responseTokenID = null;
 
            action = TrustUtil.getActionValue(RahasConstants.VERSION_05_02,
                    RahasConstants.RST_ACTION_ISSUE);
            stsClient.setAction(action);
 
            // request the security token from STS.
            Token responseToken;
             
            Policy stsPolicy = loadPolicy(stsPolicyPath);
 
            // add rampart config assertion to the ws-sec policies
            RampartConfig rampartConfig = buildRampartConfig();
            stsPolicy.addAssertion(rampartConfig);
             
            responseToken = stsClient.requestSecurityToken(null, stsEPR, stsPolicy, relyingPartyEPR);
 
            // store the obtained token in token store to be used in future
            // communication.
            TokenStorage store = TrustUtil.getTokenStore(configCtx);
            responseTokenID = responseToken.getId();
            store.add(responseToken);
 
            // print token
            System.out.println(responseToken.getToken().toString());
 
            ...
 
            //Send the token to relying party
            if (enableRelyingParty) {
                /* Invoke secured service using the obtained token */
                OMElement responseElem = null;
 
                // create service client
                ServiceClient serClient = new ServiceClient(configCtx, null);
 
                // engage modules
                serClient.engageModule("addressing");
                serClient.engageModule("rampart");
 
                // load policy of secured service
                Policy sec_policy = loadPolicy(relyingPartyPolicyPath);
 
                // add rampart config to the ws-sec policies
                sec_policy.addAssertion(rampartConfig);
 
                // set in/out security policies in client opts
                serClient.getOptions().setProperty(RampartMessageData.KEY_RAMPART_POLICY,
                        sec_policy);
 
                // Set the token id as a property in the Axis2 client scope, so that
                // this will be picked up when creating the secure message to invoke
                // the endpoint.
                serClient.getOptions().setProperty(RampartMessageData.KEY_CUSTOM_ISSUED_TOKEN,
                        responseTokenID);
 
                // set action of the Hello Service to be invoked.
                serClient.getOptions().setAction("urn:echoString");
                serClient.getOptions().setTo(new EndpointReference(relyingPartyEPR));
 
                // invoke the service
                responseElem = serClient.sendReceive(getPayload(echoRequestMsg));
                // cleanup transports
                serClient.getOptions().setCallTransportCleanup(true);
 
                System.out.println(responseElem.toString());
                 
                System.exit(0);
            }
        } catch (IOException e) {
            e.printStackTrace();
        } catch (TrustException e) {
            e.printStackTrace();
        } catch (XMLStreamException e) {
            e.printStackTrace();
        }
    }
 
    private OMElement getRSTTemplate() throws TrustException {
        OMFactory omFac = OMAbstractFactory.getOMFactory();
        OMElement element = omFac.createOMElement(SP11Constants.REQUEST_SECURITY_TOKEN_TEMPLATE);
 
        if (ClientConstants.SAML_TOKEN_TYPE_20.equals(tokenType)) {
            TrustUtil.createTokenTypeElement(RahasConstants.VERSION_05_02, element).setText(
                    RahasConstants.TOK_TYPE_SAML_20);
        } else if (ClientConstants.SAML_TOKEN_TYPE_11.equals(tokenType)) {
            TrustUtil.createTokenTypeElement(RahasConstants.VERSION_05_02, element).setText(
                    RahasConstants.TOK_TYPE_SAML_10);
        }
 
        if (ClientConstants.SUBJECT_CONFIRMATION_BEARER.equals(subjectConfirmationMethod)) {
            TrustUtil.createKeyTypeElement(RahasConstants.VERSION_05_02, element,
                    RahasConstants.KEY_TYPE_BEARER);
        } else if (ClientConstants.SUBJECT_CONFIRMATION_HOLDER_OF_KEY
                .equals(subjectConfirmationMethod)) {
            TrustUtil.createKeyTypeElement(RahasConstants.VERSION_05_02, element,
                    RahasConstants.KEY_TYPE_SYMM_KEY);
        }
 
        // request claims in the token.
        OMElement claimElement = TrustUtil.createClaims(RahasConstants.VERSION_05_02, element,claimDialect);
        // Populate the <Claims/> element with the <ClaimType/> elements
        addClaimType(claimElement, claimUris);
 
        return element;
    }
 
    private void addClaimType(OMElement parent, String[] claimUris) {
        OMElement element = null;
        // For each and every claim uri, create an <ClaimType/> elem
        for (String attr : claimUris) {
            element = parent.getOMFactory()
                    .createOMElement(
                            new QName("http://schemas.xmlsoap.org/ws/2005/05/identity",
                                    "ClaimType", "wsid"), parent);
            element.addAttribute(parent.getOMFactory().createOMAttribute("Uri", null, attr));
        }
    }
 
    private Policy loadPolicy(String policyPath) throws XMLStreamException, FileNotFoundException {
        StAXOMBuilder omBuilder = new StAXOMBuilder(policyPath);
        return PolicyEngine.getPolicy(omBuilder.getDocumentElement());
    }
 
    private RampartConfig buildRampartConfig() {
        RampartConfig rampartConfig = new RampartConfig();
        rampartConfig.setUser(username);
        rampartConfig.setEncryptionUser(encryptionUser);
        rampartConfig.setUserCertAlias(userCertAlias);
        rampartConfig.setPwCbClass(pwdCallbackClass);
 
        Properties cryptoProperties = new Properties();
        cryptoProperties.put("org.apache.ws.security.crypto.merlin.keystore.type", "JKS");
        cryptoProperties.put("org.apache.ws.security.crypto.merlin.file", keystorePath);
        cryptoProperties
                .put("org.apache.ws.security.crypto.merlin.keystore.password", keystorePwd);
 
        CryptoConfig cryptoConfig = new CryptoConfig();
        cryptoConfig.setProvider("org.apache.ws.security.components.crypto.Merlin");
        cryptoConfig.setProp(cryptoProperties);
 
        rampartConfig.setEncrCryptoConfig(cryptoConfig);
        rampartConfig.setSigCryptoConfig(cryptoConfig);
 
        return rampartConfig;
    }
 
    private OMElement getPayload(String value) {
        OMFactory factory = null;
        OMNamespace ns = null;
        OMElement elem = null;
        OMElement childElem = null;
 
        factory = OMAbstractFactory.getOMFactory();
        ns = factory.createOMNamespace("http://echo.services.core.carbon.wso2.org", "ns");
        elem = factory.createOMElement("echoString", ns);
        childElem = factory.createOMElement("in", null);
        childElem.setText(value);
        elem.addChild(childElem);
 
        return elem;
    }
     
    ...
}

public class PasswordCBHandler implements CallbackHandler{
     
    ...
     
    public void handle(Callback[] callbacks) throws IOException, UnsupportedCallbackException {
 
        readUsernamePasswordFromProperties();
         
        WSPasswordCallback pwcb = (WSPasswordCallback) callbacks[0];
        String id = pwcb.getIdentifier();
        int usage = pwcb.getUsage();
 
        if (usage == WSPasswordCallback.USERNAME_TOKEN) {
 
           if (username.equals(id)) {
               pwcb.setPassword(password);
           }
        } else if (usage == WSPasswordCallback.SIGNATURE || usage == WSPasswordCallback.DECRYPT) {
 
            if (keyAlias.equals(id)) {
                pwcb.setPassword(keyPassword);
            }
        }
    }
     
    ...
}