ws-security-stax/src/main/java/org/apache/wss4j/stax/impl/securityToken/KerberosServiceSecurityTokenImpl.java

/**
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements. See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership. The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License. You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 * KIND, either express or implied. See the License for the
 * specific language governing permissions and limitations
 * under the License.
 */
package org.apache.wss4j.stax.impl.securityToken;

import java.io.IOException;
import java.security.Key;
import java.security.Principal;
import java.security.PrivilegedActionException;
import java.util.Set;

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.LoginContext;
import javax.security.auth.login.LoginException;

import org.apache.wss4j.common.ext.WSSecurityException;
import org.apache.wss4j.common.ext.WSSecurityException.ErrorCode;
import org.apache.wss4j.common.kerberos.KerberosContextAndServiceNameCallback;
import org.apache.wss4j.common.kerberos.KerberosServiceContext;
import org.apache.wss4j.common.kerberos.KerberosServiceExceptionAction;
import org.apache.wss4j.common.kerberos.KerberosTokenDecoder;
import org.apache.wss4j.common.kerberos.KerberosTokenDecoderException;
import org.apache.wss4j.common.util.KeyUtils;
import org.apache.wss4j.stax.ext.WSInboundSecurityContext;
import org.apache.wss4j.stax.securityToken.KerberosServiceSecurityToken;
import org.apache.wss4j.stax.securityToken.WSSecurityTokenConstants;
import org.apache.xml.security.exceptions.XMLSecurityException;
import org.apache.xml.security.stax.ext.XMLSecurityConstants;
import org.apache.xml.security.stax.impl.securityToken.AbstractInboundSecurityToken;

public class KerberosServiceSecurityTokenImpl extends AbstractInboundSecurityToken implements KerberosServiceSecurityToken {

    private CallbackHandler callbackHandler;
    private byte[] binaryContent;
    private String kerberosTokenValueType;

    private KerberosTokenDecoder kerberosTokenDecoder;
    private Subject subject;
    private Principal principal;
    private byte[] sessionKey;

    public KerberosServiceSecurityTokenImpl(WSInboundSecurityContext wsInboundSecurityContext, CallbackHandler callbackHandler,
                                            byte[] binaryContent, String kerberosTokenValueType, String id,
                                            WSSecurityTokenConstants.KeyIdentifier keyIdentifier) {
        super(wsInboundSecurityContext, id, keyIdentifier, true);
        this.callbackHandler = callbackHandler;
        this.binaryContent = binaryContent;
        this.kerberosTokenValueType = kerberosTokenValueType;
    }

    @Override
    public boolean isAsymmetric() throws XMLSecurityException {
        return false;
    }

    @Override
    public WSSecurityTokenConstants.TokenType getTokenType() {
        return WSSecurityTokenConstants.KERBEROS_TOKEN;
    }

    protected byte[] getTGTSessionKey() throws WSSecurityException {
        if (sessionKey != null) {
            return sessionKey;
        }
        try {
            KerberosContextAndServiceNameCallback contextAndServiceNameCallback = new KerberosContextAndServiceNameCallback();
            callbackHandler.handle(new Callback[]{contextAndServiceNameCallback});

            if (contextAndServiceNameCallback.getContextName() == null) {
                throw new WSSecurityException(WSSecurityException.ErrorCode.FAILURE, "kerberosCallbackContextNameNotSupplied");
            }
            if (contextAndServiceNameCallback.getServiceName() == null) {
                throw new WSSecurityException(WSSecurityException.ErrorCode.FAILURE, "kerberosCallbackServiceNameNotSupplied");
            }

            LoginContext loginContext = new LoginContext(contextAndServiceNameCallback.getContextName(), callbackHandler);
            loginContext.login();

            // Get the service name to use - fall back on the principal
            this.subject = loginContext.getSubject();

            String service = contextAndServiceNameCallback.getServiceName();
            if (service == null) {
                Set<Principal> principals = subject.getPrincipals();
                if (principals.isEmpty()) {
                    throw new WSSecurityException(
                        WSSecurityException.ErrorCode.FAILURE,
                        "kerberosLoginError",
                        new Object[] {"No Client principals found after login"}
                    );
                }
                service = principals.iterator().next().getName();
            }

            KerberosServiceExceptionAction action =
                new KerberosServiceExceptionAction(binaryContent,
                                                   service,
                                                   contextAndServiceNameCallback.isUsernameServiceNameForm(),
                                                   false);
            KerberosServiceContext krbServiceCtx = null;
            try {
                krbServiceCtx = Subject.doAs(subject, action);
            } catch (PrivilegedActionException e) {
                Throwable cause = e.getCause();
                if (cause instanceof WSSecurityException) {
                    throw (WSSecurityException) cause;
                } else {
                    throw new WSSecurityException(
                        ErrorCode.FAILURE, new Exception(cause), "kerberosTicketValidationError"
                    );
                }
            }

            this.principal = krbServiceCtx.getPrincipal();

            Key key = krbServiceCtx.getSessionKey();
            if (key != null) {
                sessionKey = key.getEncoded();
            } else if (kerberosTokenDecoder != null) {
                kerberosTokenDecoder.clear();
                kerberosTokenDecoder.setToken(binaryContent);
                kerberosTokenDecoder.setSubject(subject);
                sessionKey = kerberosTokenDecoder.getSessionKey();
            }

            return sessionKey;
        } catch (LoginException | UnsupportedCallbackException | IOException e) {
            throw new WSSecurityException(WSSecurityException.ErrorCode.FAILURE, e);
        } catch (KerberosTokenDecoderException e) {
            throw new WSSecurityException(WSSecurityException.ErrorCode.INVALID_SECURITY_TOKEN, e);
        }
    }

    @Override
    protected Key getKey(String algorithmURI, XMLSecurityConstants.AlgorithmUsage algorithmUsage,
                         String correlationID) throws XMLSecurityException {

        Key key = getSecretKey().get(algorithmURI);
        if (key != null) {
            return key;
        }

        byte[] sk = getTGTSessionKey();

        key = KeyUtils.prepareSecretKey(algorithmURI, sk);
        setSecretKey(algorithmURI, key);
        return key;
    }

    public byte[] getBinaryContent() {
        return binaryContent;
    }

    public String getKerberosTokenValueType() {
        return kerberosTokenValueType;
    }

    @Override
    public Subject getSubject() throws WSSecurityException {
        return subject;
    }

    @Override
    public Principal getPrincipal() throws WSSecurityException {
        return principal;
    }
}