// 
// Decompiled by Procyon v0.6.0
// 

package com.hypixel.hytale.server.core.io.handlers.login;

import com.hypixel.hytale.server.core.auth.PlayerAuthentication;
import java.security.SecureRandom;
import com.hypixel.hytale.server.core.modules.singleplayer.SingleplayerModule;
import com.hypixel.hytale.server.core.Constants;
import com.hypixel.hytale.protocol.packets.auth.ServerAuthToken;
import com.hypixel.hytale.server.core.io.transport.QUICTransport;
import java.security.cert.X509Certificate;
import com.hypixel.hytale.protocol.io.netty.ProtocolUtil;
import com.hypixel.hytale.protocol.packets.auth.AuthGrant;
import com.hypixel.hytale.server.core.auth.ServerAuthManager;
import com.hypixel.hytale.server.core.HytaleServerConfig;
import com.hypixel.hytale.server.core.io.netty.NettyUtil;
import java.util.logging.Level;
import com.hypixel.hytale.server.core.HytaleServer;
import com.hypixel.hytale.server.core.io.PacketHandler;
import com.hypixel.hytale.protocol.packets.auth.AuthToken;
import com.hypixel.hytale.protocol.packets.connection.Disconnect;
import com.hypixel.hytale.protocol.Packet;
import com.hypixel.hytale.server.core.auth.AuthConfig;
import javax.annotation.Nullable;
import com.hypixel.hytale.server.core.io.ProtocolVersion;
import javax.annotation.Nonnull;
import io.netty.channel.Channel;
import com.hypixel.hytale.protocol.HostAddress;
import java.util.UUID;
import com.hypixel.hytale.protocol.packets.connection.ClientType;
import com.hypixel.hytale.server.core.auth.JWTValidator;
import com.hypixel.hytale.server.core.auth.SessionServiceClient;
import com.hypixel.hytale.logger.HytaleLogger;
import com.hypixel.hytale.server.core.io.handlers.GenericConnectionPacketHandler;

public abstract class HandshakeHandler extends GenericConnectionPacketHandler
{
    private static final HytaleLogger LOGGER;
    private static volatile SessionServiceClient sessionServiceClient;
    private static volatile JWTValidator jwtValidator;
    private volatile AuthState authState;
    private volatile boolean authTokenPacketReceived;
    private volatile String authenticatedUsername;
    private final ClientType clientType;
    private final String identityToken;
    private final UUID playerUuid;
    private final String username;
    private final byte[] referralData;
    private final HostAddress referralSource;
    
    public HandshakeHandler(@Nonnull final Channel channel, @Nonnull final ProtocolVersion protocolVersion, @Nonnull final String language, @Nonnull final ClientType clientType, @Nonnull final String identityToken, @Nonnull final UUID playerUuid, @Nonnull final String username, @Nullable final byte[] referralData, @Nullable final HostAddress referralSource) {
        super(channel, protocolVersion, language);
        this.authState = AuthState.REQUESTING_AUTH_GRANT;
        this.authTokenPacketReceived = false;
        this.clientType = clientType;
        this.identityToken = identityToken;
        this.playerUuid = playerUuid;
        this.username = username;
        this.referralData = referralData;
        this.referralSource = referralSource;
    }
    
    private static SessionServiceClient getSessionServiceClient() {
        if (HandshakeHandler.sessionServiceClient == null) {
            synchronized (HandshakeHandler.class) {
                if (HandshakeHandler.sessionServiceClient == null) {
                    HandshakeHandler.sessionServiceClient = new SessionServiceClient("https://sessions.hytale.com");
                }
            }
        }
        return HandshakeHandler.sessionServiceClient;
    }
    
    private static JWTValidator getJwtValidator() {
        if (HandshakeHandler.jwtValidator == null) {
            synchronized (HandshakeHandler.class) {
                if (HandshakeHandler.jwtValidator == null) {
                    HandshakeHandler.jwtValidator = new JWTValidator(getSessionServiceClient(), "https://sessions.hytale.com", AuthConfig.getServerAudience());
                }
            }
        }
        return HandshakeHandler.jwtValidator;
    }
    
    @Override
    public void accept(@Nonnull final Packet packet) {
        switch (packet.getId()) {
            case 1: {
                this.handle((Disconnect)packet);
                break;
            }
            case 12: {
                this.handle((AuthToken)packet);
                break;
            }
            default: {
                this.disconnect("Protocol error: unexpected packet " + packet.getId());
                break;
            }
        }
    }
    
    public void registered0(final PacketHandler oldHandler) {
        final HytaleServerConfig.TimeoutProfile timeouts = HytaleServer.get().getConfig().getConnectionTimeouts();
        this.enterStage("auth", timeouts.getAuth());
        final JWTValidator.IdentityTokenClaims identityClaims = getJwtValidator().validateIdentityToken(this.identityToken);
        if (identityClaims == null) {
            HandshakeHandler.LOGGER.at(Level.WARNING).log("Identity token validation failed for %s from %s", this.username, NettyUtil.formatRemoteAddress(this.channel));
            this.disconnect("Invalid or expired identity token");
            return;
        }
        final UUID tokenUuid = identityClaims.getSubjectAsUUID();
        if (tokenUuid == null || !tokenUuid.equals(this.playerUuid)) {
            HandshakeHandler.LOGGER.at(Level.WARNING).log("Identity token UUID mismatch for %s from %s (expected: %s, got: %s)", this.username, NettyUtil.formatRemoteAddress(this.channel), this.playerUuid, tokenUuid);
            this.disconnect("Invalid identity token: UUID mismatch");
            return;
        }
        final String requiredScope = (this.clientType == ClientType.Editor) ? "hytale:editor" : "hytale:client";
        if (!identityClaims.hasScope(requiredScope)) {
            HandshakeHandler.LOGGER.at(Level.WARNING).log("Identity token missing required scope for %s from %s (clientType: %s, required: %s, actual: %s)", this.username, NettyUtil.formatRemoteAddress(this.channel), this.clientType, requiredScope, identityClaims.scope);
            this.disconnect("Invalid identity token: missing " + requiredScope + " scope");
            return;
        }
        HandshakeHandler.LOGGER.at(Level.INFO).log("Identity token validated for %s (UUID: %s, scope: %s) from %s, requesting auth grant", this.username, this.playerUuid, identityClaims.scope, NettyUtil.formatRemoteAddress(this.channel));
        this.continueStage("auth:grant", timeouts.getAuthGrant(), () -> this.authState != AuthState.REQUESTING_AUTH_GRANT);
        this.requestAuthGrant();
    }
    
    private void requestAuthGrant() {
        final String serverSessionToken = ServerAuthManager.getInstance().getSessionToken();
        if (serverSessionToken == null || serverSessionToken.isEmpty()) {
            HandshakeHandler.LOGGER.at(Level.SEVERE).log("Server session token not available - cannot request auth grant");
            this.disconnect("Server authentication unavailable - please try again later");
            return;
        }
        getSessionServiceClient().requestAuthorizationGrantAsync(this.identityToken, AuthConfig.getServerAudience(), serverSessionToken).thenAccept(authGrant -> {
            if (!(!this.channel.isActive())) {
                if (authGrant == null) {
                    this.channel.eventLoop().execute(() -> this.disconnect("Failed to obtain authorization grant from session service"));
                }
                else {
                    final String serverIdentityToken = ServerAuthManager.getInstance().getIdentityToken();
                    if (serverIdentityToken == null || serverIdentityToken.isEmpty()) {
                        HandshakeHandler.LOGGER.at(Level.SEVERE).log("Server identity token not available - cannot complete mutual authentication");
                        this.channel.eventLoop().execute(() -> this.disconnect("Server authentication unavailable - please try again later"));
                    }
                    else {
                        final String finalServerIdentityToken = serverIdentityToken;
                        this.channel.eventLoop().execute(() -> {
                            if (!(!this.channel.isActive())) {
                                if (this.authState != AuthState.REQUESTING_AUTH_GRANT) {
                                    HandshakeHandler.LOGGER.at(Level.WARNING).log("State changed during auth grant request, current state: %s", this.authState);
                                }
                                else {
                                    this.clearTimeout();
                                    HandshakeHandler.LOGGER.at(Level.INFO).log("Sending AuthGrant to %s (with server identity: %s)", NettyUtil.formatRemoteAddress(this.channel), !finalServerIdentityToken.isEmpty());
                                    this.write(new AuthGrant(authGrant, finalServerIdentityToken));
                                    this.authState = AuthState.AWAITING_AUTH_TOKEN;
                                    final HytaleServerConfig.TimeoutProfile timeouts = HytaleServer.get().getConfig().getConnectionTimeouts();
                                    this.continueStage("auth:token", timeouts.getAuthToken(), () -> this.authState != AuthState.AWAITING_AUTH_TOKEN);
                                }
                            }
                        });
                    }
                }
            }
        }).exceptionally(ex -> {
            HandshakeHandler.LOGGER.at(Level.WARNING).withCause(ex).log("Error requesting auth grant");
            this.channel.eventLoop().execute(() -> this.disconnect("Authentication error: " + ex.getMessage()));
            return null;
        });
    }
    
    public void handle(@Nonnull final Disconnect packet) {
        this.disconnectReason.setClientDisconnectType(packet.type);
        HandshakeHandler.LOGGER.at(Level.INFO).log("%s (%s) at %s left with reason: %s - %s", this.playerUuid, this.username, NettyUtil.formatRemoteAddress(this.channel), packet.type.name(), packet.reason);
        ProtocolUtil.closeApplicationConnection(this.channel);
    }
    
    public void handle(@Nonnull final AuthToken packet) {
        if (this.authState != AuthState.AWAITING_AUTH_TOKEN) {
            HandshakeHandler.LOGGER.at(Level.WARNING).log("Received unexpected AuthToken packet in state %s from %s", this.authState, NettyUtil.formatRemoteAddress(this.channel));
            this.disconnect("Protocol error: unexpected AuthToken packet");
            return;
        }
        if (this.authTokenPacketReceived) {
            HandshakeHandler.LOGGER.at(Level.WARNING).log("Received duplicate AuthToken packet from %s", NettyUtil.formatRemoteAddress(this.channel));
            this.disconnect("Protocol error: duplicate AuthToken packet");
            return;
        }
        this.authTokenPacketReceived = true;
        this.authState = AuthState.PROCESSING_AUTH_TOKEN;
        this.clearTimeout();
        final String accessToken = packet.accessToken;
        if (accessToken == null || accessToken.isEmpty()) {
            HandshakeHandler.LOGGER.at(Level.WARNING).log("Received AuthToken packet with empty access token from %s", NettyUtil.formatRemoteAddress(this.channel));
            this.disconnect("Invalid access token");
            return;
        }
        final String serverAuthGrant = packet.serverAuthorizationGrant;
        final X509Certificate clientCert = this.channel.attr(QUICTransport.CLIENT_CERTIFICATE_ATTR).get();
        HandshakeHandler.LOGGER.at(Level.INFO).log("Received AuthToken from %s, validating JWT (mTLS cert present: %s, server auth grant: %s)", NettyUtil.formatRemoteAddress(this.channel), clientCert != null, serverAuthGrant != null && !serverAuthGrant.isEmpty());
        final JWTValidator.JWTClaims claims = getJwtValidator().validateToken(accessToken, clientCert);
        if (claims == null) {
            HandshakeHandler.LOGGER.at(Level.WARNING).log("JWT validation failed for %s", NettyUtil.formatRemoteAddress(this.channel));
            this.disconnect("Invalid access token");
            return;
        }
        final UUID tokenUuid = claims.getSubjectAsUUID();
        final String tokenUsername = claims.username;
        if (tokenUuid == null || !tokenUuid.equals(this.playerUuid)) {
            HandshakeHandler.LOGGER.at(Level.WARNING).log("JWT UUID mismatch for %s (expected: %s, got: %s)", NettyUtil.formatRemoteAddress(this.channel), this.playerUuid, tokenUuid);
            this.disconnect("Invalid token claims: UUID mismatch");
            return;
        }
        if (tokenUsername == null || tokenUsername.isEmpty()) {
            HandshakeHandler.LOGGER.at(Level.WARNING).log("JWT missing username for %s", NettyUtil.formatRemoteAddress(this.channel));
            this.disconnect("Invalid token claims: missing username");
            return;
        }
        if (!tokenUsername.equals(this.username)) {
            HandshakeHandler.LOGGER.at(Level.WARNING).log("JWT username mismatch for %s (expected: %s, got: %s)", NettyUtil.formatRemoteAddress(this.channel), this.username, tokenUsername);
            this.disconnect("Invalid token claims: username mismatch");
            return;
        }
        this.authenticatedUsername = tokenUsername;
        if (serverAuthGrant != null && !serverAuthGrant.isEmpty()) {
            this.authState = AuthState.EXCHANGING_SERVER_TOKEN;
            final HytaleServerConfig.TimeoutProfile timeouts = HytaleServer.get().getConfig().getConnectionTimeouts();
            this.continueStage("auth:server-exchange", timeouts.getAuthServerExchange(), () -> this.authState != AuthState.EXCHANGING_SERVER_TOKEN);
            this.exchangeServerAuthGrant(serverAuthGrant);
        }
        else {
            HandshakeHandler.LOGGER.at(Level.WARNING).log("Client did not provide server auth grant for mutual authentication");
            this.disconnect("Mutual authentication required - please update your client");
        }
    }
    
    private void exchangeServerAuthGrant(@Nonnull final String serverAuthGrant) {
        final ServerAuthManager serverAuthManager = ServerAuthManager.getInstance();
        final String serverCertFingerprint = serverAuthManager.getServerCertificateFingerprint();
        if (serverCertFingerprint == null) {
            HandshakeHandler.LOGGER.at(Level.SEVERE).log("Server certificate fingerprint not available for mutual auth");
            this.disconnect("Server authentication unavailable - please try again later");
            return;
        }
        final String serverSessionToken = serverAuthManager.getSessionToken();
        HandshakeHandler.LOGGER.at(Level.FINE).log("Server session token available: %s, identity token available: %s", serverSessionToken != null, serverAuthManager.getIdentityToken() != null);
        if (serverSessionToken == null) {
            HandshakeHandler.LOGGER.at(Level.SEVERE).log("Server session token not available for auth grant exchange");
            HandshakeHandler.LOGGER.at(Level.FINE).log("Auth mode: %s, has session token: %s, has identity token: %s", serverAuthManager.getAuthStatus(), serverAuthManager.hasSessionToken(), serverAuthManager.hasIdentityToken());
            this.disconnect("Server authentication unavailable - please try again later");
            return;
        }
        HandshakeHandler.LOGGER.at(Level.FINE).log("Using session token (first 20 chars): %s...", (serverSessionToken.length() > 20) ? serverSessionToken.substring(0, 20) : serverSessionToken);
        getSessionServiceClient().exchangeAuthGrantForTokenAsync(serverAuthGrant, serverCertFingerprint, serverSessionToken).thenAccept(serverAccessToken -> {
            if (!(!this.channel.isActive())) {
                this.channel.eventLoop().execute(() -> {
                    if (!(!this.channel.isActive())) {
                        if (this.authState != AuthState.EXCHANGING_SERVER_TOKEN) {
                            HandshakeHandler.LOGGER.at(Level.WARNING).log("State changed during server token exchange, current state: %s", this.authState);
                        }
                        else if (serverAccessToken == null) {
                            HandshakeHandler.LOGGER.at(Level.SEVERE).log("Failed to exchange server auth grant for access token");
                            this.disconnect("Server authentication failed - please try again later");
                        }
                        else {
                            final byte[] passwordChallenge = this.generatePasswordChallengeIfNeeded();
                            HandshakeHandler.LOGGER.at(Level.INFO).log("Sending ServerAuthToken to %s (with password challenge: %s)", NettyUtil.formatRemoteAddress(this.channel), passwordChallenge != null);
                            this.write(new ServerAuthToken(serverAccessToken, passwordChallenge));
                            this.completeAuthentication(passwordChallenge);
                        }
                    }
                });
            }
        }).exceptionally(ex -> {
            HandshakeHandler.LOGGER.at(Level.WARNING).withCause(ex).log("Error exchanging server auth grant");
            this.channel.eventLoop().execute(() -> {
                if (this.authState != AuthState.EXCHANGING_SERVER_TOKEN) {
                    return;
                }
                else {
                    this.disconnect("Server authentication failed - please try again later");
                    return;
                }
            });
            return null;
        });
    }
    
    private byte[] generatePasswordChallengeIfNeeded() {
        final String password = HytaleServer.get().getConfig().getPassword();
        if (password == null || password.isEmpty()) {
            return null;
        }
        if (Constants.SINGLEPLAYER) {
            final UUID ownerUuid = SingleplayerModule.getUuid();
            if (ownerUuid != null && ownerUuid.equals(this.playerUuid)) {
                return null;
            }
        }
        final byte[] challenge = new byte[32];
        new SecureRandom().nextBytes(challenge);
        return challenge;
    }
    
    private void completeAuthentication(final byte[] passwordChallenge) {
        this.auth = new PlayerAuthentication(this.playerUuid, this.authenticatedUsername);
        if (this.referralData != null) {
            this.auth.setReferralData(this.referralData);
        }
        if (this.referralSource != null) {
            this.auth.setReferralSource(this.referralSource);
        }
        this.authState = AuthState.AUTHENTICATED;
        this.clearTimeout();
        HandshakeHandler.LOGGER.at(Level.INFO).log("Mutual authentication complete for %s (%s) from %s", this.authenticatedUsername, this.playerUuid, NettyUtil.formatRemoteAddress(this.channel));
        this.onAuthenticated(passwordChallenge);
    }
    
    protected abstract void onAuthenticated(final byte[] p0);
    
    static {
        LOGGER = HytaleLogger.forEnclosingClass();
    }
    
    private enum AuthState
    {
        REQUESTING_AUTH_GRANT, 
        AWAITING_AUTH_TOKEN, 
        PROCESSING_AUTH_TOKEN, 
        EXCHANGING_SERVER_TOKEN, 
        AUTHENTICATED;
    }
}