Auth Service (OAuth / SSO)

// ─── JAVA BLUEPRINT ──────────────────────────────────────────────────────────
import java.util.*;
import java.security.MessageDigest;
import java.nio.charset.StandardCharsets;
import javax.crypto.Mac;
import javax.crypto.spec.SecretKeySpec;

import java.util.concurrent.ConcurrentHashMap;

class User {
    private final String userId;
    private final String username;
    private final String passwordHash;
    private final Set<String> roles;
    private final Map<String, String> attributes;

    public User(String userId, String username, String passwordHash, Set<String> roles, Map<String, String> attributes) {
        this.userId = userId;
        this.username = username;
        this.passwordHash = passwordHash;
        this.roles = Collections.unmodifiableSet(new HashSet<>(roles));
        this.attributes = Collections.unmodifiableMap(new HashMap<>(attributes));
    }

    public String getUserId() { return userId; }
    public String getUsername() { return username; }
    public String getPasswordHash() { return passwordHash; }
    public Set<String> getRoles() { return roles; }
    public Map<String, String> getAttributes() { return attributes; }
}

class OAuthClient {
    private final String clientId;
    private final String clientSecretHash;
    private final Set<String> allowedScopes;
    private final Set<String> redirectUris;

    public OAuthClient(String clientId, String clientSecretHash, Set<String> allowedScopes, Set<String> redirectUris) {
        this.clientId = clientId;
        this.clientSecretHash = clientSecretHash;
        this.allowedScopes = Collections.unmodifiableSet(new HashSet<>(allowedScopes));
        this.redirectUris = Collections.unmodifiableSet(new HashSet<>(redirectUris));
    }

    public String getClientId() { return clientId; }
    public String getClientSecretHash() { return clientSecretHash; }
    public Set<String> getAllowedScopes() { return allowedScopes; }
    public Set<String> getRedirectUris() { return redirectUris; }
}

class AuthorizationCode {
    private final String code;
    private final String clientId;
    private final String userId;
    private final String redirectUri;
    private final Set<String> scopes;
    private final long expiresAt;
    private final String codeChallenge;
    private final String codeChallengeMethod;

    public AuthorizationCode(String code, String clientId, String userId, String redirectUri, Set<String> scopes, long ttlMs, String codeChallenge, String codeChallengeMethod) {
        this.code = code;
        this.clientId = clientId;
        this.userId = userId;
        this.redirectUri = redirectUri;
        this.scopes = Collections.unmodifiableSet(new HashSet<>(scopes));
        this.expiresAt = System.currentTimeMillis() + ttlMs;
        this.codeChallenge = codeChallenge;
        this.codeChallengeMethod = codeChallengeMethod;
    }

    public boolean isExpired() { return System.currentTimeMillis() > expiresAt; }
    public String getCode() { return code; }
    public String getClientId() { return clientId; }
    public String getUserId() { return userId; }
    public String getRedirectUri() { return redirectUri; }
    public Set<String> getScopes() { return scopes; }
    public String getCodeChallenge() { return codeChallenge; }
    public String getCodeChallengeMethod() { return codeChallengeMethod; }
}

class TokenPair {
    private final String accessToken;
    private final String refreshToken;
    private final long expiresAt;

    public TokenPair(String accessToken, String refreshToken, long ttlMs) {
        this.accessToken = accessToken;
        this.refreshToken = refreshToken;
        this.expiresAt = System.currentTimeMillis() + ttlMs;
    }

    public String getAccessToken() { return accessToken; }
    public String getRefreshToken() { return refreshToken; }
    public long getExpiresAt() { return expiresAt; }
}

class RefreshTokenMetadata {
    final String token;
    final String parentToken;
    final String userId;
    final String clientId;
    final Set<String> scopes;
    volatile boolean isUsed;
    volatile boolean isRevoked;

    public RefreshTokenMetadata(String token, String parentToken, String userId, String clientId, Set<String> scopes) {
        this.token = token;
        this.parentToken = parentToken;
        this.userId = userId;
        this.clientId = clientId;
        this.scopes = Collections.unmodifiableSet(new HashSet<>(scopes));
        this.isUsed = false;
        this.isRevoked = false;
    }
}

class AuthService {
    private final Map<String, User> userDb = new ConcurrentHashMap<>();
    private final Map<String, OAuthClient> clientDb = new ConcurrentHashMap<>();
    private final Map<String, AuthorizationCode> codeDb = new ConcurrentHashMap<>();
    private final Map<String, RefreshTokenMetadata> refreshTokenDb = new ConcurrentHashMap<>();
    private final Map<String, Set<String>> rbacPolicy = new ConcurrentHashMap<>();
    private final String jwtSigningSecret = "oauth-secret-key-signature-hmac-sha-256";

    public synchronized void registerUser(User user) { userDb.put(user.getUserId(), user); }
    public synchronized void registerClient(OAuthClient client) { clientDb.put(client.getClientId(), client); }
    public synchronized void addRolePermission(String role, String permission) {
        rbacPolicy.computeIfAbsent(role, k -> ConcurrentHashMap.newKeySet()).add(permission);
    }

    public synchronized String issueAuthorizationCode(String clientId, String userId, String redirectUri, Set<String> scopes, String challenge, String method) {
        OAuthClient client = clientDb.get(clientId);
        if (client == null) throw new IllegalArgumentException("Client not found");
        if (!client.getRedirectUris().contains(redirectUri)) throw new IllegalArgumentException("Invalid redirect URI");

        String code = UUID.randomUUID().toString();
        AuthorizationCode authCode = new AuthorizationCode(code, clientId, userId, redirectUri, scopes, 60000, challenge, method);
        codeDb.put(code, authCode);
        return code;
    }

    public synchronized TokenPair exchangeCode(String code, String clientId, String clientSecret, String redirectUri, String verifier) throws Exception {
        AuthorizationCode authCode = codeDb.remove(code);
        if (authCode == null) throw new SecurityException("Invalid/Used auth code");
        if (authCode.isExpired()) throw new SecurityException("Auth code expired");

        OAuthClient client = clientDb.get(clientId);
        if (client == null) throw new IllegalArgumentException("Client not found");
        
        if (!client.getClientSecretHash().equals(hashSecret(clientSecret))) {
            throw new SecurityException("Client credentials check failed");
        }
        if (!authCode.getRedirectUri().equals(redirectUri)) {
            throw new SecurityException("Redirect URI mismatch");
        }

        if (authCode.getCodeChallenge() != null) {
            if (verifier == null) throw new SecurityException("PKCE verifier required");
            if (!verifyPKCE(authCode.getCodeChallenge(), authCode.getCodeChallengeMethod(), verifier)) {
                throw new SecurityException("PKCE verification failed");
            }
        }

        return generateTokenPair(authCode.getUserId(), clientId, authCode.getScopes(), null);
    }

    public synchronized TokenPair refreshTokens(String refreshToken) throws Exception {
        RefreshTokenMetadata meta = refreshTokenDb.get(refreshToken);
        if (meta == null) throw new SecurityException("Refresh token does not exist");

        if (meta.isUsed || meta.isRevoked) {
            revokeFamilyTree(meta.token);
            throw new SecurityException("Token reuse alert! Full token family has been revoked.");
        }

        meta.isUsed = true;
        return generateTokenPair(meta.userId, meta.clientId, meta.scopes, meta.token);
    }

    private synchronized TokenPair generateTokenPair(String userId, String clientId, Set<String> scopes, String parentToken) throws Exception {
        User user = userDb.get(userId);
        if (user == null) throw new IllegalArgumentException("User not found");

        String accessToken = generateMockJWT(user, clientId, scopes);
        String refreshToken = UUID.randomUUID().toString();
        
        RefreshTokenMetadata meta = new RefreshTokenMetadata(refreshToken, parentToken, userId, clientId, scopes);
        refreshTokenDb.put(refreshToken, meta);

        return new TokenPair(accessToken, refreshToken, 900000);
    }

    private synchronized void revokeFamilyTree(String rootToken) {
        RefreshTokenMetadata meta = refreshTokenDb.get(rootToken);
        if (meta != null) {
            meta.isRevoked = true;
            for (RefreshTokenMetadata m : refreshTokenDb.values()) {
                if (rootToken.equals(m.parentToken)) {
                    revokeFamilyTree(m.token);
                }
            }
        }
    }

    private String generateMockJWT(User user, String clientId, Set<String> scopes) throws Exception {
        long exp = System.currentTimeMillis() + 900000;
        String header = Base64.getUrlEncoder().withoutPadding().encodeToString("{\"alg\":\"HS256\",\"typ\":\"JWT\"}".getBytes(StandardCharsets.UTF_8));
        String rolesStr = "[" + String.join(",", user.getRoles().stream().map(r -> "\"" + r + "\"").toArray(String[]::new)) + "]";
        String payloadJson = String.format("{\"sub\":\"%s\",\"aud\":\"%s\",\"roles\":%s,\"exp\":%d}", user.getUserId(), clientId, rolesStr, exp);
        String payload = Base64.getUrlEncoder().withoutPadding().encodeToString(payloadJson.getBytes(StandardCharsets.UTF_8));
        
        Mac mac = Mac.getInstance("HmacSHA256");
        mac.init(new SecretKeySpec(jwtSigningSecret.getBytes(StandardCharsets.UTF_8), "HmacSHA256"));
        byte[] sigBytes = mac.doFinal((header + "." + payload).getBytes(StandardCharsets.UTF_8));
        String sig = Base64.getUrlEncoder().withoutPadding().encodeToString(sigBytes);

        return header + "." + payload + "." + sig;
    }

    private boolean verifyPKCE(String challenge, String method, String verifier) throws Exception {
        if ("plain".equalsIgnoreCase(method)) {
            return challenge.equals(verifier);
        } else if ("S256".equalsIgnoreCase(method)) {
            MessageDigest digest = MessageDigest.getInstance("SHA-256");
            byte[] hash = digest.digest(verifier.getBytes(StandardCharsets.UTF_8));
            String calculated = Base64.getUrlEncoder().withoutPadding().encodeToString(hash);
            return challenge.equals(calculated);
        }
        return false;
    }

    private String hashSecret(String secret) throws Exception {
        MessageDigest digest = MessageDigest.getInstance("SHA-256");
        byte[] hash = digest.digest(secret.getBytes(StandardCharsets.UTF_8));
        return Base64.getEncoder().encodeToString(hash);
    }

    public synchronized boolean checkAccess(String accessToken, String resourceDept, int requiredClearance) {
        try {
            String[] parts = accessToken.split(".");
            if (parts.length < 3) return false;
            
            String payload = new String(Base64.getUrlDecoder().decode(parts[1]), StandardCharsets.UTF_8);
            String userId = null;
            int startSub = payload.indexOf("\"sub\":\"");
            if (startSub != -1) {
                int endSub = payload.indexOf("\"", startSub + 7);
                userId = payload.substring(startSub + 7, endSub);
            }
            
            if (userId == null) return false;
            User user = userDb.get(userId);
            if (user == null) return false;
            
            // RBAC Check
            boolean hasRole = user.getRoles().contains("Administrator") || user.getRoles().contains("Developer");
            System.out.println("    [RBAC Check] User has roles " + user.getRoles() + " -> " + (hasRole ? "PASSED" : "FAILED"));
            if (!hasRole) return false;
            
            // ABAC Check
            String userDept = user.getAttributes().getOrDefault("department", "");
            int userClearance = Integer.parseInt(user.getAttributes().getOrDefault("clearanceLevel", "0"));
            
            boolean deptMatch = userDept.equalsIgnoreCase(resourceDept);
            boolean clearanceMatch = userClearance >= requiredClearance;
            
            System.out.println("    [ABAC Check] Dept matches (" + userDept + " vs " + resourceDept + ") -> " + (deptMatch ? "PASSED" : "FAILED"));
            System.out.println("    [ABAC Check] Clearance level " + userClearance + " >= " + requiredClearance + " -> " + (clearanceMatch ? "PASSED" : "FAILED"));
            
            return deptMatch && clearanceMatch;
        } catch (Exception e) {
            System.out.println("    [Access Control] Token validation error: " + e.getMessage());
            return false;
        }
    }
}

public class Main {
    public static void main(String[] args) {
        try {
            System.out.println("=== STARTING AUTH & SSO SERVICE SIMULATION ===");
            AuthService authService = new AuthService();

            // 1. Setup Identities & Clients
            Set<String> devRoles = new HashSet<>(Arrays.asList("Developer"));
            Map<String, String> devAttrs = new HashMap<>();
            devAttrs.put("department", "Engineering");
            devAttrs.put("clearanceLevel", "2");
            User bob = new User("usr-bob", "Bob Developer", "hashed_password_123", devRoles, devAttrs);
            authService.registerUser(bob);

            MessageDigest digest = MessageDigest.getInstance("SHA-256");
            String clientSecret = "client-secret-xyz-789";
            String secretHash = Base64.getEncoder().encodeToString(digest.digest(clientSecret.getBytes(StandardCharsets.UTF_8)));
            
            OAuthClient client = new OAuthClient("client-portal", secretHash, new HashSet<>(Arrays.asList("openid", "profile")), new HashSet<>(Arrays.asList("https://client.com/callback")));
            authService.registerClient(client);

            System.out.println("[Init] Registered User: Bob (Developer, Engineering, Clearance: 2)");
            System.out.println("[Init] Registered OAuth Client: client-portal");

            // 2. PKCE Authorization Flow S256
            String verifier = "my-super-secure-pkce-verifier-string-12345-bob-verifier";
            byte[] hash = digest.digest(verifier.getBytes(StandardCharsets.UTF_8));
            String challenge = Base64.getUrlEncoder().withoutPadding().encodeToString(hash);

            System.out.println("\n[Flow 1] Requesting short-lived authorization code with PKCE S256...");
            String code = authService.issueAuthorizationCode("client-portal", "usr-bob", "https://client.com/callback", new HashSet<>(Arrays.asList("openid")), challenge, "S256");
            System.out.println("  -> Authorization Code issued: " + code);

            System.out.println("\n[Flow 2] Exchanging authorization code with verifier...");
            TokenPair pair1 = authService.exchangeCode(code, "client-portal", clientSecret, "https://client.com/callback", verifier);
            System.out.println("  -> Token exchange successful!");
            System.out.println("  -> Access Token (JWT snippet): " + pair1.getAccessToken().substring(0, 35) + "...");
            System.out.println("  -> Refresh Token (Opaque): " + pair1.getRefreshToken());

            // 3. RBAC + ABAC Authorization
            System.out.println("\n[Flow 3] Verifying access to protected resource: /api/secured-dataset (Engineering dept, Clearance >= 2)");
            boolean allowed = authService.checkAccess(pair1.getAccessToken(), "Engineering", 2);
            System.out.println("  -> Authorization decision: ACCESS " + (allowed ? "GRANTED" : "DENIED"));

            System.out.println("\n[Flow 3b] Verifying access to protected resource: /api/finance-records (Finance dept, Clearance >= 3)");
            boolean allowedFinance = authService.checkAccess(pair1.getAccessToken(), "Finance", 3);
            System.out.println("  -> Authorization decision: ACCESS " + (allowedFinance ? "GRANTED" : "DENIED"));

            // 4. Refresh Token Rotation (RTR)
            System.out.println("\n[Flow 4] Refreshing session using token: " + pair1.getRefreshToken());
            TokenPair pair2 = authService.refreshTokens(pair1.getRefreshToken());
            System.out.println("  -> Session rotated successfully!");
            System.out.println("  -> New Access Token (JWT snippet): " + pair2.getAccessToken().substring(0, 35) + "...");
            System.out.println("  -> New Refresh Token (Opaque): " + pair2.getRefreshToken());

            // 5. Token Replay Attack Countermeasures
            System.out.println("\n[Flow 5] Simulating Token Replay Attack using the rotated (invalidated) refresh token...");
            try {
                authService.refreshTokens(pair1.getRefreshToken());
            } catch (SecurityException e) {
                System.out.println("  -> SECURITY ALERT CAUGHT: " + e.getMessage());
            }

            System.out.println("\n[Flow 5b] Verifying if cascade revocation successfully invalidated the new active session...");
            try {
                authService.refreshTokens(pair2.getRefreshToken());
            } catch (SecurityException e) {
                System.out.println("  -> SUCCESS: Current session also revoked successfully! Message: " + e.getMessage());
            }

            System.out.println("\n=== SIMULATION COMPLETED SUCCESSFULLY ===");
        } catch (Exception e) {
            e.printStackTrace();
        }
    }
}