package com.chequer.axboot.admin.utils;

import com.chequer.axboot.admin.domain.user.SessionUser;
import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.ObjectMapper;
import org.springframework.security.core.authority.SimpleGrantedAuthority;

import javax.crypto.Mac;
import javax.crypto.spec.SecretKeySpec;
import javax.xml.bind.DatatypeConverter;
import java.io.ByteArrayInputStream;
import java.io.IOException;
import java.security.InvalidKeyException;
import java.security.NoSuchAlgorithmException;
import java.util.Arrays;

public final class JWTSessionHandler {

    private static final String HMAC_ALGO = "HmacSHA256";
    private static final String SEPARATOR = ".";
    private static final String SEPARATOR_SPLITTER = "\\.";

    private final Mac hmac;

    public JWTSessionHandler(byte[] secretKey) {
        try {
            hmac = Mac.getInstance(HMAC_ALGO);
            hmac.init(new SecretKeySpec(secretKey, HMAC_ALGO));
        } catch (NoSuchAlgorithmException | InvalidKeyException e) {
            throw new IllegalStateException("failed to initialize HMAC: " + e.getMessage(), e);
        }
    }

    public SessionUser parseUserFromToken(String token) {
        final String[] parts = token.split(SEPARATOR_SPLITTER);
        if (parts.length == 2 && parts[0].length() > 0 && parts[1].length() > 0) {
            try {
                final byte[] userBytes = fromBase64(parts[0]);
                final byte[] hash = fromBase64(parts[1]);

                boolean validHash = Arrays.equals(createHmac(userBytes), hash);
                if (validHash) {
                    return fromJSON(userBytes);
                }
            } catch (Exception e) {
                e.printStackTrace();
            }
        }
        return null;
    }

    public String createTokenForUser(SessionUser user) {
        byte[] userBytes = toJSON(user);
        byte[] hash = createHmac(userBytes);
        final StringBuilder sb = new StringBuilder(170);
        sb.append(toBase64(userBytes));
        sb.append(SEPARATOR);
        sb.append(toBase64(hash));
        return sb.toString();
    }

    private SessionUser fromJSON(final byte[] userBytes) {
        try {
            ObjectMapper objectMapper = new ObjectMapper();
            objectMapper.registerSubtypes(SimpleGrantedAuthority.class);
            return objectMapper.readValue(new ByteArrayInputStream(userBytes), SessionUser.class);
        } catch (IOException e) {
            throw new IllegalStateException(e);
        }
    }

    private byte[] toJSON(SessionUser user) {
        try {
            ObjectMapper objectMapper = new ObjectMapper();
            objectMapper.registerSubtypes(SimpleGrantedAuthority.class);
            return objectMapper.writeValueAsBytes(user);
        } catch (JsonProcessingException e) {
            throw new IllegalStateException(e);
        }
    }

    private String toBase64(byte[] content) {
        return DatatypeConverter.printBase64Binary(content);
    }

    private byte[] fromBase64(String content) {
        return DatatypeConverter.parseBase64Binary(content);
    }

    private synchronized byte[] createHmac(byte[] content) {
        return hmac.doFinal(content);
    }
}

@@ -198,24 +198,71 @@
 jectMapper;%0A
+import org.apache.commons.codec.binary.Base64;%0A
 import org.s
@@ -394,49 +394,8 @@
 ec;%0A
-import javax.xml.bind.DatatypeConverter;%0A
 impo
@@ -2984,43 +2984,40 @@
 urn 
-DatatypeConverter.printBase64Binary
+Base64.encodeBase64URLSafeString
 (con
@@ -3097,43 +3097,27 @@
 urn 
-DatatypeConverter.pars
+Base64.decod
 eBase64
-Binary
 (con

Replace "javax.xml.bind.DatatypeConverter" with "org.apache.commons.codec.binary.Base64", which is used for Base64 encoding/decoding. "DatatypeConverter" caused a serious problem in "WildFly 9". The problem is that the session stored in the cookie can not be restored properly.