package io.quarkus.jwt.test;

import static org.junit.jupiter.api.Assertions.assertNotNull;
import static org.junit.jupiter.api.Assertions.assertTrue;

import java.security.KeyPair;
import java.security.KeyPairGenerator;
import java.security.Permission;
import java.security.Principal;
import java.security.PrivateKey;
import java.security.PublicKey;
import java.security.interfaces.RSAPublicKey;
import java.util.HashSet;
import java.util.Iterator;
import java.util.Set;

import org.junit.jupiter.api.Test;
import org.wildfly.security.auth.principal.NamePrincipal;
import org.wildfly.security.auth.realm.token.TokenSecurityRealm;
import org.wildfly.security.auth.server.RealmIdentity;
import org.wildfly.security.auth.server.SecurityDomain;
import org.wildfly.security.auth.server.SecurityIdentity;
import org.wildfly.security.authz.Attributes;
import org.wildfly.security.authz.AuthorizationIdentity;
import org.wildfly.security.authz.PermissionMappable;
import org.wildfly.security.authz.PermissionMapper;
import org.wildfly.security.authz.RoleDecoder;
import org.wildfly.security.authz.Roles;
import org.wildfly.security.authz.SimpleAttributesEntry;
import org.wildfly.security.evidence.BearerTokenEvidence;
import org.wildfly.security.permission.PermissionVerifier;

import io.quarkus.smallrye.jwt.runtime.auth.MpJwtValidator;
import io.smallrye.jwt.auth.principal.JWTAuthContextInfo;

/**
 * Validate usage of the bearer token based realm
 */
public class TokenRealmUnitTest {

    @Test
    public void testTokenRealm() throws Exception {
        KeyPair keyPair = KeyPairGenerator.getInstance("RSA").generateKeyPair();
        PublicKey pk1 = keyPair.getPublic();
        PrivateKey pk1Priv = keyPair.getPrivate();
        JWTAuthContextInfo contextInfo = new JWTAuthContextInfo();
        contextInfo.setSignerKey((RSAPublicKey) pk1);
        contextInfo.setIssuedBy("https://server.example.com");
        MpJwtValidator jwtValidator = new MpJwtValidator(contextInfo);

        TokenSecurityRealm tokenRealm = TokenSecurityRealm.builder()
                .claimToPrincipal(this::mpJwtLogic)
                .validator(jwtValidator)
                .build();

        String jwt = TokenUtils.generateTokenString("/Token1.json", pk1Priv, "testTokenRealm");
        BearerTokenEvidence tokenEvidence = new BearerTokenEvidence(jwt);
        tokenRealm.getRealmIdentity(tokenEvidence);
        RealmIdentity identity = tokenRealm.getRealmIdentity(tokenEvidence);
        assertNotNull(identity);
        assertTrue(identity.exists());
        AuthorizationIdentity authz = identity.getAuthorizationIdentity();
        // TODO add proper assertion
        //System.out.println(authz.getAttributes().keySet());
    }

    @Test
    public void testSecurityDomain() throws Exception {
        KeyPair keyPair = KeyPairGenerator.getInstance("RSA").generateKeyPair();
        PublicKey pk1 = keyPair.getPublic();
        PrivateKey pk1Priv = keyPair.getPrivate();
        JWTAuthContextInfo contextInfo = new JWTAuthContextInfo();
        contextInfo.setSignerKey((RSAPublicKey) pk1);
        contextInfo.setIssuedBy("https://server.example.com");
        MpJwtValidator jwtValidator = new MpJwtValidator(contextInfo);
        TokenSecurityRealm tokenRealm = TokenSecurityRealm.builder()
                .claimToPrincipal(this::mpJwtLogic)
                .validator(jwtValidator)
                .build();

        SecurityDomain securityDomain = SecurityDomain.builder()
                .addRealm("MP-JWT", tokenRealm)
                .setRoleDecoder(new RoleDecoder() {
                    @Override
                    public Roles decodeRoles(AuthorizationIdentity authorizationIdentity) {
                        SimpleAttributesEntry groups = (SimpleAttributesEntry) authorizationIdentity.getAttributes()
                                .get("groups");
                        Set<String> roles = new HashSet<>(groups);
                        return new Roles() {
                            @Override
                            public boolean contains(String roleName) {
                                return roles.contains(roleName);
                            }

                            @Override
                            public Iterator<String> iterator() {
                                return roles.iterator();
                            }
                        };
                    }
                })
                .build()
                .setPermissionMapper(new PermissionMapper() {
                    @Override
                    public PermissionVerifier mapPermissions(PermissionMappable permissionMappable, Roles roles) {
                        return new PermissionVerifier() {
                            @Override
                            public boolean implies(Permission permission) {
                                return true;
                            }
                        };
                    }
                })
                .build();

        String jwt = TokenUtils.generateTokenString("/Token1.json", pk1Priv, "testTokenRealm");
        BearerTokenEvidence tokenEvidence = new BearerTokenEvidence(jwt);
        SecurityIdentity securityIdentity = securityDomain.authenticate(tokenEvidence);
        // TODO add proper assertion
        //System.out.println(securityIdentity.getAttributes().keySet());
    }

    private Principal mpJwtLogic(Attributes claims) {
        String pn = claims.getFirst("upn");
        if (pn == null) {
            pn = claims.getFirst("preferred_name");
        }
        if (pn == null) {
            pn = claims.getFirst("sub");
        }
        return new NamePrincipal(pn);
    }
}

@@ -206,32 +206,79 @@
 yPairGenerator;%0A
+import java.security.NoSuchAlgorithmException;%0A
 import java.secu
@@ -290,24 +290,24 @@
 Permission;%0A
-
 import java.
@@ -1616,44 +1616,8 @@
 r = 
-KeyPairGenerator.getInstance(%22RSA%22).
 gene
@@ -2820,44 +2820,8 @@
 r = 
-KeyPairGenerator.getInstance(%22RSA%22).
 gene
@@ -5606,32 +5606,32 @@
 ub%22);%0A        %7D%0A
-
         return n
@@ -5657,11 +5657,288 @@
 );%0A    %7D
+%0A%0A    private KeyPair generateKeyPair() throws NoSuchAlgorithmException %7B%0A        KeyPairGenerator generator = KeyPairGenerator.getInstance(%22RSA%22);%0A        generator.initialize(2048); // because that's the minimal accepted size%0A        return generator.generateKeyPair();%0A    %7D
 %0A%7D%0A

Fix test failing due to default RSA keypair size too small (depending on the JDK)