package com.novell.oauth.samples;

import java.io.ByteArrayOutputStream;
import java.io.FileInputStream;
import java.io.IOException;
import java.net.MalformedURLException;
import java.security.spec.RSAPrivateCrtKeySpec;
import java.util.List;
import java.util.Scanner;

import org.apache.commons.io.IOUtils;
import org.jose4j.jwa.AlgorithmConstraints;
import org.jose4j.jwe.JsonWebEncryption;
import org.jose4j.jwk.JsonWebKey;
import org.jose4j.jwk.JsonWebKeySet;
import org.jose4j.jwk.RsaJsonWebKey;
import org.jose4j.jws.JsonWebSignature;
import org.jose4j.lang.JoseException;


public class IDTokenDecoder {

	public static List<JsonWebKey> readKeySet(String jwksFileName) throws IOException, JoseException{
		String jwks;
		try(ByteArrayOutputStream baos = new ByteArrayOutputStream();
				FileInputStream fin = new FileInputStream(jwksFileName)){
				IOUtils.copy(fin, baos);
				jwks = baos.toString("UTF-8");
			}
			JsonWebKeySet jsonWebKeySet = new JsonWebKeySet(jwks);
			return jsonWebKeySet.getJsonWebKeys();
	}
	
	public static void displayIDTokenWhenSignAlgNONE(String jwt) throws JoseException{
		
		JsonWebSignature jws = new JsonWebSignature();
	    jws.setCompactSerialization(jwt);
	    jws.setAlgorithmConstraints(AlgorithmConstraints.ALLOW_ONLY_NONE);
	    String payload = jws.getPayload();
	    System.out.println("idtoken:");
	    System.out.println(payload);

	}
	
	public static void displayIdTokenIfSignatureMatches(String jwt, String jwksFileName) throws JoseException, MalformedURLException, IOException{
		
		List<JsonWebKey> jsonWebKeys = readKeySet(jwksFileName);
		JsonWebKey jsonWebkey = jsonWebKeys.get(0);
		if(jsonWebkey == null){
			System.out.println("json web key is null" );
			return;
		}
		JsonWebSignature jws = new JsonWebSignature();
	    jws.setKey(jsonWebkey.getKey());;
	    jws.setCompactSerialization(jwt);
		 
	    if(false == jws.verifySignature()){
	    	System.out.println("signature mismatch");
	    	return;
	    }
	    System.out.println("Signature verified successfully");
	    String payload = jws.getPayload();
	    System.out.println("idtoken:");
	    System.out.println(payload);
	    
	}
	public static String getIDTokenIfDecryptionSuccessful(String jwt,String jwksFileName) throws IOException, JoseException{
		List<JsonWebKey> jsonWebKeys = readKeySet(jwksFileName);
		JsonWebKey jsonWebkey = jsonWebKeys.get(0);
		if(jsonWebkey == null){
			System.out.println("json web key is null" );
			return null;
		}
		
		JsonWebEncryption jwe = new JsonWebEncryption();
        jwe.setCompactSerialization(jwt);
        if(jsonWebkey instanceof RsaJsonWebKey){
        	RsaJsonWebKey rsa = (RsaJsonWebKey) jsonWebkey;
        	jwe.setKey(rsa.getPrivateKey());
        }else
        {
        	jwe.setKey(jsonWebkey.getKey());
        }

        String decrString = jwe.getPlaintextString();
        return decrString;
 	}
	public static void displayMainOptions()
	{
		System.out.println("option a: plain");
		System.out.println("option b:verify signature");
		System.out.println("option c:decryptjwt");
		System.out.println("option d:DecryptAndVerifySignature");
		System.out.println("option e:Exit");
	}
	public static void main (String [] args){
		Scanner scanner = null;
		try
		{
			scanner = new Scanner(System.in);
			
			System.out.print("Enter your jwt token:");
			String jwt = scanner.next().trim();
			displayMainOptions();
			System.out.println("Enter your option:");
			String input = scanner.next().trim();
			switch(input.charAt(0))
			{
				case 'a':
							{
								displayIDTokenWhenSignAlgNONE(jwt);
								break;
							}
				case 'b' :
							{
									System.out.print("Enter IDP's JWKS file path:");
									String jwksFilePath = scanner.next();
									displayIdTokenIfSignatureMatches(jwt,jwksFilePath);
									break;
							}
				case 'c' :
							{
								System.out.print("Enter the client's JWKS file path:");
								String jwksFilePath = scanner.next();
								String idToken = getIDTokenIfDecryptionSuccessful(jwt,jwksFilePath);
								System.out.println(idToken);
								break;
								
							}
				case 'd' :
							{
								System.out.print("Enter the client's JWKS file path:");
								String clientFilePath = scanner.next();
								System.out.print("Enter IDP's JWKS file path:");
								String idpFilePath = scanner.next();
								String idToken = getIDTokenIfDecryptionSuccessful(jwt,clientFilePath);
								System.out.println("decrypted jwt is :" + idToken);
								displayIdTokenIfSignatureMatches(idToken,idpFilePath);	
								break;
							}
				case 'e':
							{
								return;
							}
				default:
							{
								System.out.println("invalid option");
							}
							
			}
		}
		catch(Exception e){
			System.out.println(e.getMessage());
			e.printStackTrace();
		}
	}
}
