package com.wmc.demo.teiid;

import java.security.Principal;
import java.security.PrivilegedActionException;
import java.security.PrivilegedExceptionAction;
import java.util.HashSet;
import java.util.Set;

import javax.security.auth.Subject;
import javax.security.auth.kerberos.KerberosPrincipal;
import javax.security.auth.login.LoginContext;
import javax.security.auth.login.LoginException;

import org.ietf.jgss.GSSContext;
import org.ietf.jgss.GSSCredential;
import org.ietf.jgss.GSSManager;
import org.teiid.dqp.service.GSSResult;
import org.teiid.logging.LogConstants;
import org.teiid.logging.LogManager;
import org.teiid.net.socket.AuthenticationType;
import org.teiid.services.SessionServiceImpl;

import com.sun.security.jgss.GSSUtil;

public class KerbSessionServiceImpl extends SessionServiceImpl {

	final Subject serviceSubject;
	final LoginContext lc;

	/**
	 * Implementation of SessionService that defaults to GSS/Kerberos authentication
	 * @param servicePrincipalName the SPN registered for this service (e.g. postgres/my-server.domain.com@ad-domain.COM)
	 * @param kerberosDomain
	 * @throws LoginException in case of problem reading Kerberos keytab file
	 */
	public KerbSessionServiceImpl(String servicePrincipalName, String kerberosDomain) throws LoginException {
		this.setAuthenticationType(AuthenticationType.GSS);
		this.setSecurityDomain(kerberosDomain);
		LogManager.logInfo(LogConstants.CTX_SECURITY, "Establishing Kerberos LoginContext for " + servicePrincipalName);
		Set<Principal> princ = new HashSet<Principal>(1);
		princ.add(new KerberosPrincipal(servicePrincipalName));
		Subject sub = new Subject(false, princ, new HashSet<Object>(), new HashSet<Object>());
		this.lc = new LoginContext("teiid-server", sub);
		this.lc.login();
		this.serviceSubject = lc.getSubject();
	}

	@Override
	protected GSSResult neogitiateGssLogin(final String securityDomain,	final byte[] serviceTicket) throws LoginException {
		try {
			return Subject.doAs(serviceSubject, new PrivilegedExceptionAction<GSSResult>() {
				public GSSResult run() throws Exception {
					GSSContext context = GSSManager.getInstance().createContext((GSSCredential) null);
					byte[] outKey = context.acceptSecContext(serviceTicket, 0, serviceTicket.length);
					GSSCredential delCreds = context.getCredDelegState() ? context.getDelegCred() : null;
					GSSResult r = new GSSResult(outKey, context.isEstablished(), delCreds);
					String userName = context.getSrcName().toString();
					r.setUserName(userName);
					Subject clientSubject = GSSUtil.createSubject(context.getSrcName(), delCreds);
					r.setSecurityContext(new KerbSecurityContext(clientSubject, userName, securityDomain));
					context.dispose();
					LogManager.logInfo(LogConstants.CTX_SECURITY, "Kerberos login accepted from " + userName);
					return r;
				}
			});
		} catch (PrivilegedActionException e) {
			LoginException le = new LoginException("Cannot read client kerberos token due to: " + e.toString());
			le.initCause(e);
			throw le;
		}
	}
}