digital-rmb-backend/src/main/java/com/yau/digitalrmb/security/application/JwtTokenService.java

package com.yau.digitalrmb.security.application;

import com.yau.digitalrmb.identity.domain.UserRole;
import com.yau.digitalrmb.identity.infrastructure.persistence.entity.UserEntity;
import com.yau.digitalrmb.security.config.SecurityProperties;
import org.springframework.security.oauth2.jwt.Jwt;
import org.springframework.security.oauth2.jwt.JwtEncoder;
import org.springframework.security.oauth2.jwt.JwtEncoderParameters;
import org.springframework.security.oauth2.jwt.JwtClaimsSet;
import org.springframework.security.oauth2.jose.jws.MacAlgorithm;
import org.springframework.security.oauth2.jwt.JwsHeader;
import org.springframework.stereotype.Service;

import java.time.Instant;
import java.util.List;
import java.util.Set;
import java.util.Collections;
import java.util.stream.Collectors;

@Service
public class JwtTokenService {
    private final JwtEncoder jwtEncoder;
    private final SecurityProperties properties;

    public JwtTokenService(JwtEncoder jwtEncoder, SecurityProperties properties) {
        this.jwtEncoder = jwtEncoder;
        this.properties = properties;
    }

    public Token issueFor(String username) {
        Instant issuedAt = Instant.now();
        Instant expiresAt = issuedAt.plus(properties.getJwt().getAccessTokenTtl());
        JwtClaimsSet claims = JwtClaimsSet.builder().subject(username).issuedAt(issuedAt).expiresAt(expiresAt)
                .claim("roles", Collections.singletonList("ROLE_ADMIN")).build();
        Jwt jwt = jwtEncoder.encode(JwtEncoderParameters.from(JwsHeader.with(MacAlgorithm.HS256).build(), claims));
        return new Token(jwt.getTokenValue(), properties.getJwt().getAccessTokenTtl().getSeconds());
    }

    public Token issueFor(String userId, String account, Set<String> roles) {
        return issueFor(userId, null, null, account, roles);
    }

    public Token issueFor(String userId, String account, int roleId) {
        return issueFor(userId, account, Collections.singleton(UserRole.fromId(roleId).getAuthority()));
    }

    public Token issueFor(UserEntity user) {
        Instant issuedAt = Instant.now();
        Instant expiresAt = issuedAt.plus(properties.getJwt().getAccessTokenTtl());
        List<String> authorities = Collections.singletonList(
                "ROLE_" + UserRole.fromId(user.getRoleId()).getAuthority());
        JwtClaimsSet.Builder claims = JwtClaimsSet.builder()
                .subject(user.getUserId())
                .issuedAt(issuedAt)
                .expiresAt(expiresAt)
                .claim("userId", user.getUserId())
                .claim("preferred_username", user.getStudentId())
                .claim("studentId", user.getStudentId())
                .claim("roleId", user.getRoleId())
                .claim("roles", authorities);
        claimIfPresent(claims, "userName", user.getUserName());
        claimIfPresent(claims, "classId", user.getClassId());
        claimIfPresent(claims, "className", user.getClassName());
        claimIfPresent(claims, "phone", user.getPhone());
        claimIfPresent(claims, "schoolId", user.getSchoolId());
        claimIfPresent(claims, "schoolName", user.getSchoolName());
        claimIfPresent(claims, "authorizeTime", user.getAuthorizeTime());
        claimIfPresent(claims, "authorizeEndTime", user.getAuthorizeEndTime());
        claimIfPresent(claims, "createTime", user.getCreateTime());
        claimIfPresent(claims, "isDeleted", user.getIsDeleted());
        claimIfPresent(claims, "zyUserId", user.getZyUserId());
        Jwt jwt = jwtEncoder.encode(JwtEncoderParameters.from(JwsHeader.with(MacAlgorithm.HS256).build(),
                claims.build()));
        return new Token(jwt.getTokenValue(), properties.getJwt().getAccessTokenTtl().getSeconds());
    }

    public Token issueFor(String userId, Long schoolId, Long classId, String account, Set<String> roles) {
        Instant issuedAt = Instant.now();
        Instant expiresAt = issuedAt.plus(properties.getJwt().getAccessTokenTtl());
        List<String> authorities = roles.stream().sorted().map(role -> "ROLE_" + role).collect(Collectors.toList());
        JwtClaimsSet.Builder claims = JwtClaimsSet.builder()
                .subject(userId)
                .issuedAt(issuedAt)
                .expiresAt(expiresAt)
                .claim("userId", userId)
                .claim("preferred_username", account)
                .claim("roles", authorities);
        if (schoolId != null) {
            claims.claim("schoolId", schoolId);
        }
        if (classId != null) {
            claims.claim("classId", classId);
        }
        Jwt jwt = jwtEncoder.encode(JwtEncoderParameters.from(JwsHeader.with(MacAlgorithm.HS256).build(),
                claims.build()));
        return new Token(jwt.getTokenValue(), properties.getJwt().getAccessTokenTtl().getSeconds());
    }

    private static void claimIfPresent(JwtClaimsSet.Builder claims, String name, Object value) {
        if (value != null) {
            claims.claim(name, value.toString());
        }
    }

    public static class Token {
        private final String accessToken;
        private final long expiresIn;
        public Token(String accessToken, long expiresIn) { this.accessToken = accessToken; this.expiresIn = expiresIn; }
        public String accessToken() { return accessToken; }
        public long expiresIn() { return expiresIn; }
    }
}