/*
 * Copyright (c) 2016 Network New Technologies Inc.
 *
 * Licensed 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 com.networknt.client.oauth;

import com.networknt.client.ClientConfig;
import com.networknt.client.Http2Client;
import com.networknt.config.Config;
import com.networknt.monad.Result;
import io.undertow.Handlers;
import io.undertow.Undertow;
import io.undertow.util.Headers;
import org.jose4j.jws.AlgorithmIdentifiers;
import org.jose4j.jws.JsonWebSignature;
import org.jose4j.jwt.JwtClaims;
import org.jose4j.jwt.NumericDate;
import org.jose4j.jwt.consumer.InvalidJwtException;
import org.jose4j.jwt.consumer.JwtConsumer;
import org.jose4j.jwt.consumer.JwtConsumerBuilder;
import org.jose4j.lang.JoseException;
import org.junit.*;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.nio.ByteBuffer;
import java.security.KeyStore;
import java.security.PrivateKey;
import java.security.interfaces.RSAPrivateKey;
import java.util.*;
import java.util.regex.Pattern;

/**
 * This is the tests for OauthHelper and it doesn't need live light-oauth2
 * server up and running.
 *
 */
public class OauthHelperTest {
    static final Logger logger = LoggerFactory.getLogger(OauthHelperTest.class);
    static final String token = "eyJraWQiOiIxMDAiLCJhbGciOiJSUzI1NiJ9.eyJpc3MiOiJ1cm46Y29tOm5ldHdvcmtudDpvYXV0aDI6djEiLCJhdWQiOiJ1cm46Y29tLm5ldHdvcmtudCIsImV4cCI6MTc5MDAzNTcwOSwianRpIjoiSTJnSmdBSHN6NzJEV2JWdUFMdUU2QSIsImlhdCI6MTQ3NDY3NTcwOSwibmJmIjoxNDc0Njc1NTg5LCJ2ZXJzaW9uIjoiMS4wIiwidXNlcl9pZCI6InN0ZXZlIiwidXNlcl90eXBlIjoiRU1QTE9ZRUUiLCJjbGllbnRfaWQiOiJmN2Q0MjM0OC1jNjQ3LTRlZmItYTUyZC00YzU3ODc0MjFlNzIiLCJzY29wZSI6WyJ3cml0ZTpwZXRzIiwicmVhZDpwZXRzIl19.mue6eh70kGS3Nt2BCYz7ViqwO7lh_4JSFwcHYdJMY6VfgKTHhsIGKq2uEDt3zwT56JFAePwAxENMGUTGvgceVneQzyfQsJeVGbqw55E9IfM_uSM-YcHwTfR7eSLExN4pbqzVDI353sSOvXxA98ZtJlUZKgXNE1Ngun3XFORCRIB_eH8B0FY_nT_D1Dq2WJrR-re-fbR6_va95vwoUdCofLRa4IpDfXXx19ZlAtfiVO44nw6CS8O87eGfAm7rCMZIzkWlCOFWjNHnCeRsh7CVdEH34LF-B48beiG5lM7h4N12-EME8_VDefgMjZ8eqs1ICvJMxdIut58oYbdnkwTjkA";
    static Undertow server = null;

    @BeforeClass
    public static void setUp() {
        if(server == null) {
            logger.info("starting server");
            server = Undertow.builder()
                    .addHttpListener(8887, "localhost")
                    .setHandler(Handlers.header(Handlers.path()
                                    .addPrefixPath("/oauth2/key", (exchange) -> {
                                        exchange.getResponseHeaders().put(Headers.CONTENT_TYPE, "application/json");
                                        exchange.getResponseSender().send("OK");
                                    })
                                    .addPrefixPath("/oauth2/introspection", (exchange) -> {
                                        exchange.getResponseHeaders().put(Headers.CONTENT_TYPE, "application/json");
                                        Map<String, Object> map = new HashMap<>();
                                        map.put("active", true);
                                        map.put("client_id", "bb8293f6-ceef-4e7a-90c8-1492e97df19f");
                                        map.put("token_type", "access_token");
                                        map.put("scope", "openid profile");
                                        map.put("sub", "cn=odicuser,dc=example,dc=com");
                                        map.put("exp", "86400");
                                        map.put("iat", "1506513918");
                                        map.put("iss", "https://wa.example.com");
                                        exchange.getResponseSender().send(Config.getInstance().getMapper().writeValueAsString(map));
                                    })
                                    .addPrefixPath("/oauth2/deref", (exchange) -> {
                                        exchange.getResponseHeaders().put(Headers.CONTENT_TYPE, "application/json");
                                        exchange.getResponseSender().send(token);
                                    })
                                    .addPrefixPath("/oauth2/token", (exchange) -> {
                                        // create a token that expired in 5 seconds.
                                        Map<String, Object> map = new HashMap<>();
                                        String token = getJwt(5);
                                        map.put("access_token", token);
                                        map.put("token_type", "Bearer");
                                        map.put("expires_in", 5);
                                        exchange.getResponseHeaders().put(Headers.CONTENT_TYPE, "application/json");
                                        exchange.getResponseSender().send(ByteBuffer.wrap(
                                                Config.getInstance().getMapper().writeValueAsBytes(map)));
                                    }),
                            Headers.SERVER_STRING, "U-tow"))
                    .build();
            server.start();
        }
    }

    @AfterClass
    public static void tearDown() throws Exception {
        if(server != null) {
            try {
                Thread.sleep(1000);
            } catch (InterruptedException ignored) {
            }
            server.stop();
            System.out.println("The server is stopped.");
            try {
                Thread.sleep(1000);
            } catch (InterruptedException ignored) {
            }
        }
    }

    private static boolean isTokenExpired(String authorization) {
        boolean expired = false;
        String jwt = getJwtFromAuthorization(authorization);
        if(jwt != null) {
            JwtConsumer consumer = new JwtConsumerBuilder()
                    .setDisableRequireSignature()
                    .setSkipSignatureVerification()
                    .build();

            try {
                consumer.processToClaims(jwt);
            } catch (InvalidJwtException e) {
                if(e.hasExpired()) expired = true;
            }
        }
        return expired;
    }

    private static String getJwt(int expiredInSeconds) throws Exception {
        JwtClaims claims = getTestClaims();
        claims.setExpirationTime(NumericDate.fromMilliseconds(System.currentTimeMillis() + expiredInSeconds * 1000));
        return getJwt(claims);
    }

    private static JwtClaims getTestClaims() {
        JwtClaims claims = new JwtClaims();
        claims.setIssuer("urn:com:networknt:oauth2:v1");
        claims.setAudience("urn:com.networknt");
        claims.setExpirationTimeMinutesInTheFuture(10);
        claims.setGeneratedJwtId(); // a unique identifier for the token
        claims.setIssuedAtToNow();  // when the token was issued/created (now)
        claims.setNotBeforeMinutesInThePast(2); // time before which the token is not yet valid (2 minutes ago)
        claims.setClaim("version", "1.0");

        claims.setClaim("user_id", "steve");
        claims.setClaim("user_type", "EMPLOYEE");
        claims.setClaim("client_id", "aaaaaaaa-1234-1234-1234-bbbbbbbb");
        List<String> scope = Arrays.asList("api.r", "api.w");
        claims.setStringListClaim("scope", scope); // multi-valued claims work too and will end up as a JSON array
        return claims;
    }

    public static String getJwtFromAuthorization(String authorization) {
        String jwt = null;
        if(authorization != null) {
            String[] parts = authorization.split(" ");
            if (parts.length == 2) {
                String scheme = parts[0];
                String credentials = parts[1];
                Pattern pattern = Pattern.compile("^Bearer$", Pattern.CASE_INSENSITIVE);
                if (pattern.matcher(scheme).matches()) {
                    jwt = credentials;
                }
            }
        }
        return jwt;
    }

    public static String getJwt(JwtClaims claims) throws JoseException {
        String jwt;

        RSAPrivateKey privateKey = (RSAPrivateKey) getPrivateKey(
                "/config/primary.jks", "password", "selfsigned");

        // A JWT is a JWS and/or a JWE with JSON claims as the payload.
        // In this example it is a JWS nested inside a JWE
        // So we first create a JsonWebSignature object.
        JsonWebSignature jws = new JsonWebSignature();

        // The payload of the JWS is JSON content of the JWT Claims
        jws.setPayload(claims.toJson());

        // The JWT is signed using the sender's private key
        jws.setKey(privateKey);
        jws.setKeyIdHeaderValue("100");

        // Set the signature algorithm on the JWT/JWS that will integrity protect the claims
        jws.setAlgorithmHeaderValue(AlgorithmIdentifiers.RSA_USING_SHA256);

        // Sign the JWS and produce the compact serialization, which will be the inner JWT/JWS
        // representation, which is a string consisting of three dot ('.') separated
        // base64url-encoded parts in the form Header.Payload.Signature
        jwt = jws.getCompactSerialization();
        return jwt;
    }

    private static PrivateKey getPrivateKey(String filename, String password, String key) {
        PrivateKey privateKey = null;

        try {
            KeyStore keystore = KeyStore.getInstance("JKS");
            keystore.load(Http2Client.class.getResourceAsStream(filename),
                    password.toCharArray());

            privateKey = (PrivateKey) keystore.getKey(key,
                    password.toCharArray());
        } catch (Exception e) {
            logger.error("Exception:", e);
        }

        if (privateKey == null) {
            logger.error("Failed to retrieve private key from keystore");
        }

        return privateKey;
    }

    @Test
    @Ignore
    public void testGetTokenResult() throws Exception {
        AuthorizationCodeRequest tokenRequest = new AuthorizationCodeRequest();
        tokenRequest.setClientId("test_client");
        tokenRequest.setClientSecret("test_secret");
        tokenRequest.setGrantType(ClientConfig.AUTHORIZATION_CODE);
        List<String> list = new ArrayList<>();
        list.add("test.r");
        list.add("test.w");
        tokenRequest.setScope(list);
        tokenRequest.setServerUrl("http://localhost:8887");
        tokenRequest.setEnableHttp2(true);
        tokenRequest.setUri("/oauth2/token");
        tokenRequest.setRedirectUri("https://localhost:8443/authorize");
        tokenRequest.setAuthCode("test_code");

        Result<TokenResponse> result = OauthHelper.getTokenResult(tokenRequest);
        Assert.assertTrue(result.isSuccess());
        TokenResponse tokenResponse = result.getResult();
        System.out.println("tokenResponse = " + tokenResponse);
    }

    @Test
    @Deprecated
    @Ignore
    public void testGetToken() throws Exception {
        AuthorizationCodeRequest tokenRequest = new AuthorizationCodeRequest();
        tokenRequest.setClientId("test_client");
        tokenRequest.setClientSecret("test_secret");
        tokenRequest.setGrantType(ClientConfig.AUTHORIZATION_CODE);
        List<String> list = new ArrayList<>();
        list.add("test.r");
        list.add("test.w");
        tokenRequest.setScope(list);
        tokenRequest.setServerUrl("http://localhost:8887");
        tokenRequest.setEnableHttp2(true);
        tokenRequest.setUri("/oauth2/token");
        tokenRequest.setRedirectUri("https://localhost:8443/authorize");
        tokenRequest.setAuthCode("test_code");

        Result<TokenResponse> result = OauthHelper.getTokenResult(tokenRequest);
        Assert.assertTrue(result.isSuccess());
        Assert.assertNotNull(result.getResult());
        System.out.println("tokenResponse = " + result.getResult());
    }

    @Test
    public void testGetSignKey() throws Exception {
        SignKeyRequest request = new SignKeyRequest("100");
        request.setClientId("test_client");
        request.setClientSecret("test_secret");
        request.setServerUrl("http://localhost:8887");
        request.setUri("/oauth2/key");
        request.setEnableHttp2(true);

        String key = OauthHelper.getKey(request);
        System.out.println("key = " + key);
    }

    @Test
    @Ignore
    public void testGetTokenKey() throws Exception {
        TokenKeyRequest request = new TokenKeyRequest("100");
        request.setClientId("test_client");
        request.setClientSecret("test_secret");
        request.setServerUrl("http://localhost:8887");
        request.setUri("/oauth2/key");
        request.setEnableHttp2(true);

        String key = OauthHelper.getKey(request);
        System.out.println("key = " + key);
    }

    @Test
    public void testDerefToken() throws Exception {
        DerefRequest derefRequest = new DerefRequest("xW-IFk66TKGlxLWDEaDkvg");
        derefRequest.setClientId("test_client");
        derefRequest.setClientSecret("test_secret");
        derefRequest.setServerUrl("http://localhost:8887");
        derefRequest.setUri("/oauth2/deref");
        derefRequest.setEnableHttp2(true);

        String jwt = OauthHelper.derefToken(derefRequest);
        System.out.println("jwt = " + jwt);
    }

    @Test
    @Ignore
    public void testTokenIntrospection() throws Exception {
        TokenIntrospectionRequest request = new TokenIntrospectionRequest("Simple Web Token");
        request.setClientId("test_client");
        request.setClientSecret("test_secret");
        request.setServerUrl("http://localhost:8887");
        request.setUri("/oauth2/introspection");
        request.setEnableHttp2(true);

        Result<String> introspection = OauthHelper.getIntrospection("dummy", request);
        System.out.println("introspection = " + introspection.getResult());
    }

}
