package io.digdag.server;

import com.google.common.base.Optional;
import com.google.common.collect.ImmutableMap;
import com.google.inject.Inject;
import io.digdag.client.api.RestApiKey;
import io.digdag.client.config.Config;
import io.jsonwebtoken.Claims;
import io.jsonwebtoken.JwsHeader;
import io.jsonwebtoken.JwtException;
import io.jsonwebtoken.Jwts;
import io.jsonwebtoken.SignatureException;
import io.jsonwebtoken.SigningKeyResolver;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import javax.crypto.spec.SecretKeySpec;
import javax.ws.rs.container.ContainerRequestContext;

import java.security.Key;
import java.util.Map;

public class JwtAuthenticator
    implements Authenticator
{
    private static final Logger logger = LoggerFactory.getLogger(JwtAuthenticator.class);

    private final Map<String, UserConfig> userMap;
    private final boolean allowPublicAccess;

    @Inject
    public JwtAuthenticator(Config systemConfig)
    {
        Optional<RestApiKey> apiKey = systemConfig.getOptional("server.apikey", RestApiKey.class);

        if (apiKey.isPresent()) {
            UserConfig user = UserConfig.builder()
                .siteId(0)
                .isAdmin(true)
                .apiKey(apiKey.get())
                .build();
            this.userMap = ImmutableMap.of(user.getApiKey().getIdString(), user);
            this.allowPublicAccess = false;
        }
        else {
            this.userMap = ImmutableMap.of();
            this.allowPublicAccess = true;
        }
    }

    @Override
    public Result authenticate(ContainerRequestContext requestContext)
    {
        int siteId;
        boolean admin;

        String auth = requestContext.getHeaderString("Authorization");
        if (auth == null) {
            if (allowPublicAccess) {
                // OK
                siteId = 0;
                admin = true;
            }
            else {
                return Result.reject("Authorization is required");
            }
        }
        else {
            String[] typeData = auth.split(" ", 2);
            if (typeData.length != 2) {
                return Result.reject("Invalid authorization header");
            }
            if (!typeData[0].equals("Bearer")) {
                return Result.reject("Invalid authorization header");
            }
            String token = typeData[1];
            try {
                String subject = Jwts.parser().setSigningKeyResolver(new SigningKeyResolver() {
                    @Override
                    public Key resolveSigningKey(JwsHeader header, Claims claims)
                    {
                        Object keyType = header.get("knd");
                        if (keyType == null || !"ps1".equals(keyType)) {
                            throw new SignatureException("Invalid key type");
                        }
                        UserConfig user = userMap.get(claims.getSubject());
                        if (user == null) {
                            throw new SignatureException("Invalid subject");
                        }
                        return new SecretKeySpec(user.getApiKey().getSecret(), header.getAlgorithm());
                    }

                    @Override
                    public Key resolveSigningKey(JwsHeader header, String plaintext)
                    {
                        throw new SignatureException("Plain text JWT authorization header is not allowed");
                    }
                })
                .parseClaimsJws(token)
                .getBody()
                .getSubject();

                UserConfig user = userMap.get(subject);
                if (user == null) {
                    throw new SignatureException("Invalid subject");
                }

                siteId = user.getSiteId();
                admin = user.isAdmin();
            }
            catch (JwtException ex) {
                logger.trace("Authentication failed: ", ex);
                return Result.reject("Authorization failed");
            }
        }

        return Result.builder()
                .siteId(siteId)
                .isAdmin(admin)
                .build();
    }
}

@@ -219,16 +219,102 @@
 Config;%0A
+import io.digdag.client.config.ConfigFactory;%0Aimport io.digdag.spi.AuthenticatedUser;%0A
 import i
@@ -336,16 +336,16 @@
 Claims;%0A
-
 import i
@@ -872,24 +872,60 @@
 or.class);%0A%0A
+    private final ConfigFactory cf;%0A
     private 
@@ -1068,16 +1068,40 @@
 emConfig
+, final ConfigFactory cf
 )%0A    %7B%0A
@@ -1639,32 +1639,32 @@
 cAccess = true;%0A
-
         %7D%0A    %7D%0A
@@ -1653,24 +1653,46 @@
 ;%0A        %7D%0A
+        this.cf = cf;%0A
     %7D%0A%0A    @
@@ -4232,80 +4232,468 @@
 ult.
-builder()%0A                .siteId(siteId)%0A                .isAdmin(admin
+accept(%0A                createAuthenticatedUser(siteId, admin, requestContext),%0A                () -%3E ImmutableMap.of());%0A    %7D%0A%0A    private AuthenticatedUser createAuthenticatedUser(final int siteId, final boolean admin, final ContainerRequestContext requestContext)%0A    %7B%0A        return AuthenticatedUser.builder()%0A                .siteId(siteId)%0A                .isAdmin(admin)%0A                .userInfo(cf.create())%0A                .userContext(cf.create()
 )%0A  

Refactor and modify server.JwtAuthenticator to use spi.AuthenticatedUser