/*
 * Copyright 2018 Kerbaya Software
 * 
 * This file is part of headlamp. 
 * 
 * headlamp is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 * 
 * headlamp 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 General Public License for more details.
 * 
 * You should have received a copy of the GNU General Public License
 * along with headlamp.  If not, see <http://www.gnu.org/licenses/>.
 */
package com.kerbaya.headlamp;

import java.io.IOException;

import javax.security.auth.login.LoginException;

import org.bouncycastle.util.encoders.Base64;
import org.ietf.jgss.GSSException;
import org.ietf.jgss.Oid;
import org.junit.Assert;
import org.junit.Test;

import com.kerbaya.headlamp.kerberos.Krb5TokenFactory;

public class Krb5TokenFactoryTest extends KdcTest
{
	private static byte[] establish(byte[] initToken) 
			throws GSSException, LoginException
	{
		System.out.println("-> " + Base64.toBase64String(initToken));
		Krb5TokenFactory tt = Krb5TokenFactory.create(initToken);
		initToken = tt.getGssInit();
		final byte[] response;
		if (initToken == null)
		{
			response = tt.getGssContinue();
		}
		else
		{
			System.out.println("=> " + Base64.toBase64String(initToken));
			try (TestLoginContext serviceCtx = new TestLoginContext(
							JaasConfigFile.SERVER_CONFIG_NAME);
					TestGssContext gssCtx = new TestGssContext(serviceCtx))
			{
				 
				final byte[] krb5Response = gssCtx.acceptSecContext(initToken);
				Assert.assertTrue(gssCtx.isEstablished());
				System.out.println(
						"<- " 
						+ Base64.toBase64String(krb5Response));
				response = tt.wrapGssAccept(krb5Response);
			}
		}
		if (response != null)
		{
			System.out.println(
					"<= " + Base64.toBase64String(response));
		}
		return response;
	}
	
	private static void testMech(final Oid mechOid) 
			throws LoginException, GSSException
	{
		try (TestLoginContext loginCtx = new TestLoginContext(
						JaasConfigFile.USER_1_CONFIG_NAME);
				TestGssContext gssCtx = new TestGssContext(
						loginCtx, TestUser.SERVER, mechOid);)
		{
			byte[] token = new byte[0];
			while ((token = gssCtx.initSecContext(token)) != null)
			{
				token = establish(token);
			}
			Assert.assertTrue(gssCtx.isEstablished());
		}
	}
	
	@Test
	public void testKrb5() throws LoginException, GSSException
	{
		testMech(TestUtil.KRB5_MECH_OID);
	}
	
	@Test
	public void testSpnego() throws LoginException, GSSException
	{
		testMech(TestUtil.SPNEGO_MECH_OID);
	}
	
	@Test
	public void unsupportedMech() throws LoginException, GSSException, IOException
	{
		/*
		 * KRB5 not supported
		 */
		Integer majorCode = null;
		try
		{
			establish(TokenFactory.createSpnegoUnsupportedMechInit());
			Assert.fail();
		}
		catch (GSSException e)
		{
			majorCode = e.getMajor();
		}
		Assert.assertEquals(Integer.valueOf(GSSException.BAD_MECH), majorCode);
	}
	
	@Test
	public void testSpnego2Stage() throws IOException, LoginException, GSSException
	{
		/*
		 * Phase 1: supports KRB5, but not preferred
		 */
		byte[] token1 = establish(TokenFactory.createSpnegoStage1Init());
		TokenFactory.assertSpnegoStage1Response(token1);
		
		try (TestLoginContext userCtx = new TestLoginContext(
						JaasConfigFile.USER_1_CONFIG_NAME);
				TestGssContext gssCtx = new TestGssContext(
						userCtx, TestUser.SERVER, TestUtil.KRB5_MECH_OID))
		{
			byte[] unwrappedToken2 = gssCtx.initSecContext();
			
			byte[] token2 = TokenFactory.createSpnegoStage2Init(unwrappedToken2);
			/*
			 * Phase 2: KRB5 token
			 */
			byte[] token3 = establish(token2);
			final byte[] unwrappedToken3 = TokenFactory.assertSpnegoStage2Response(token3);
			Assert.assertNull(gssCtx.initSecContext(unwrappedToken3));
			Assert.assertTrue(gssCtx.isEstablished());
		}
		
	}
}
