Flow 34 – Composite Key Rotation (HYBRID_KEM_COMBINE)

/*
 * Copyright 2025 ANKATech Solutions Inc
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 *
 * SPDX-License-Identifier: Apache-2.0
 */
package co.ankatech.ankasecure.sdk.examples;

import co.ankatech.ankasecure.sdk.AuthenticatedSdk;
import co.ankatech.ankasecure.sdk.exception.AnkaSecureSdkException;
import co.ankatech.ankasecure.sdk.model.*;
import co.ankatech.ankasecure.sdk.util.FileIO;

import java.nio.file.Files;
import java.nio.file.Path;
import java.util.Properties;

import static co.ankatech.ankasecure.sdk.examples.ExampleUtil.*;

/**
 * Scenario 34 — Composite Key Rotation (HYBRID_KEM_COMBINE).
 *
 * <p>Demonstrates quantum-resistant key rotation maintaining mode and algorithms:</p>
 * <ol>
 *   <li>Generate original composite key (X25519 + ML-KEM-768)</li>
 *   <li>Rotate to successor with same mode and algorithms</li>
 *   <li>Verify rotation chain (originalKey.nextKid &rarr; successor)</li>
 *   <li>Verify old key status = ROTATED</li>
 *   <li>Demonstrate encryption uses successor transparently</li>
 * </ol>
 *
 * <p><strong>Key Validation Rules:</strong></p>
 * <ul>
 *   <li>Same mode required (HYBRID_KEM_COMBINE &rarr; HYBRID_KEM_COMBINE)</li>
 *   <li>Same algorithms required (X25519 + ML-KEM-768 &rarr; X25519 + ML-KEM-768)</li>
 *   <li>Atomic rotation (all components together)</li>
 *   <li>Mode change blocked (HYBRID_KEM &rarr; DUALSIGN)</li>
 *   <li>Algorithm change blocked (X25519 &rarr; RSA-3072)</li>
 *   <li>Downgrade blocked (COMPOSITE &rarr; SIMPLE)</li>
 * </ul>
 *
 * @author ANKATech Solutions Inc.
 * @since 3.0.0
 * @see ExampleUtil
 * @see AuthenticatedSdk
 */
public final class ExampleScenario34 {

    private static final Path TEMP_DIR = Path.of("temp_files");

    /** No instantiation &mdash; this class only exposes {@link #main(String[])}. */
    private ExampleScenario34() { }

    /**
     * Runs the composite key rotation scenario.
     *
     * <p>Loads CLI properties, authenticates against ANKASecure&copy;,
     * and delegates to the scenario logic. On any unrecoverable error
     * the JVM terminates via
     * {@link ExampleUtil#fatal(String, Throwable)}.</p>
     *
     * @param args command-line arguments (ignored)
     */
    public static void main(String[] args) {
        System.out.println("===== SCENARIO 34: COMPOSITE KEY ROTATION =====");
        System.out.println("Purpose: Rotate hybrid quantum-resistant keys");
        System.out.println("Pattern: Generate Composite → Rotate → Verify Chain");
        System.out.println();

        try {
            ensureTempDir(TEMP_DIR);
            Properties props = loadProperties();
            AuthenticatedSdk sdk = authenticate(props);

            runScenario(sdk);

            System.out.println("===== SCENARIO 34 END =====");
        } catch (Exception ex) {
            fatal("Scenario 34 failed", ex);
        }
    }

    /* ====================================================================== */
    private static void runScenario(AuthenticatedSdk sdk) throws Exception {

        // PHASE 1: Generate Original Composite Key
        System.out.println("╔═══════════════════════════════════════════╗");
        System.out.println("║  PHASE 1: GENERATE ORIGINAL COMPOSITE    ║");
        System.out.println("╚═══════════════════════════════════════════╝");

        String originalKid = generateOriginalCompositeKey(sdk);

        // PHASE 2: Rotate to Successor
        System.out.println("╔═══════════════════════════════════════════╗");
        System.out.println("║  PHASE 2: ROTATE TO SUCCESSOR            ║");
        System.out.println("╚═══════════════════════════════════════════╝");

        String successorKid = rotateToSuccessor(sdk, originalKid);

        // PHASE 3: Verify Rotation Chain
        System.out.println("╔═══════════════════════════════════════════╗");
        System.out.println("║  PHASE 3: VERIFY ROTATION CHAIN          ║");
        System.out.println("╚═══════════════════════════════════════════╝");

        verifyRotationChain(sdk, originalKid, successorKid);

        // PHASE 4: Demonstrate Transparent Encryption
        System.out.println("╔═══════════════════════════════════════════╗");
        System.out.println("║  PHASE 4: TRANSPARENT ENCRYPTION         ║");
        System.out.println("╚═══════════════════════════════════════════╝");

        demonstrateTransparentEncryption(sdk, originalKid, successorKid);

        // Final Summary
        System.out.println();
        System.out.println("✅ SCENARIO 34 SUCCESSFUL");
        System.out.println("   - Original key: " + originalKid + " (ROTATED)");
        System.out.println("   - Successor key: " + successorKid + " (ACTIVE)");
        System.out.println("   - Rotation chain: verified");
        System.out.println("   - Quantum resistance: maintained");
    }

    private static String generateOriginalCompositeKey(AuthenticatedSdk sdk) throws Exception {
        String kid = "sc34_hybrid_v1_" + System.currentTimeMillis();

        System.out.println("[1/3] Building composite key specification...");
        System.out.println("      KID: " + kid);
        System.out.println("      Mode: HYBRID_KEM_COMBINE");
        System.out.println("      Components:");
        System.out.println("        • Classical: X25519 (Curve25519 ECDH, Level 3)");
        System.out.println("        • PQC: ML-KEM-768 (NIST FIPS 203, Level 3)");
        System.out.println("      KDF: HKDF-SHA256");

        GenerateCompositeKeySpec spec = new GenerateCompositeKeySpec()
                .setKid(kid)
                .setMode(GenerateCompositeKeySpec.Mode.HYBRID_KEM_COMBINE)
                .addComponent(ComponentSpec.classical(ClassicalCompositeAlgorithm.X25519))
                .addComponent(ComponentSpec.pqc(PqcCompositeAlgorithm.ML_KEM_768))
                .setKdf(Kdf.HKDF_SHA256)
                .setMaxUsageLimit(1000000)
                .setExportable(true);

        System.out.println();
        System.out.println("[2/3] Generating composite key on server...");

        KeyGenerationSummarySpec result = sdk.generateCompositeKey(spec);

        System.out.println("      ✅ Composite key generated");
        System.out.println("      Algorithm: " + result.getAlg());
        System.out.println("      Status: " + result.getStatus());
        System.out.println("      Usage count: " + result.getUsageCount());
        System.out.println();

        return kid;
    }

    private static String rotateToSuccessor(AuthenticatedSdk sdk, String originalKid) throws Exception {
        String successorKid = "sc34_hybrid_v2_" + System.currentTimeMillis();

        System.out.println("[3/5] Building successor specification...");
        System.out.println("      Original KID: " + originalKid);
        System.out.println("      Successor KID: " + successorKid);
        System.out.println("      ⚠️ IMPORTANT:");
        System.out.println("         • Mode MUST match (HYBRID_KEM_COMBINE)");
        System.out.println("         • Algorithms MUST match (X25519 + ML-KEM-768)");
        System.out.println("         • Cannot change to DUALSIGN");
        System.out.println("         • Cannot downgrade to SIMPLE");

        GenerateCompositeKeySpec successorSpec = new GenerateCompositeKeySpec()
                .setKid(successorKid)
                .setMode(GenerateCompositeKeySpec.Mode.HYBRID_KEM_COMBINE)  // Same mode
                .addComponent(ComponentSpec.classical(ClassicalCompositeAlgorithm.X25519))  // Same algorithm
                .addComponent(ComponentSpec.pqc(PqcCompositeAlgorithm.ML_KEM_768))  // Same algorithm
                .setKdf(Kdf.HKDF_SHA256)
                .setMaxUsageLimit(2000000);  // Can change lifecycle settings

        System.out.println();
        System.out.println("[4/5] Executing composite key rotation...");
        System.out.println("      Calling: sdk.rotateCompositeKey()");

        KeyGenerationSummarySpec successor = sdk.rotateCompositeKey(originalKid, successorSpec);

        System.out.println("      ✅ Rotation successful");
        System.out.println("      Successor KID: " + successor.getKid());
        System.out.println("      Successor Status: " + successor.getStatus());
        System.out.println();

        return successorKid;
    }

    private static void verifyRotationChain(AuthenticatedSdk sdk, String originalKid, String successorKid)
            throws Exception {

        System.out.println("[5/7] Retrieving original key metadata...");

        // Note: In v3.0.0, getKeyMetadata() returns KeyMetadata (custom model)
        // which may not have all fields populated from KeySummaryResponse
        // For full metadata with rotation chain, use exportKey() instead
        ExportedKeySpec originalMeta = sdk.exportKey(originalKid);

        System.out.println("      Original key:");
        System.out.println("         • KID: " + originalMeta.getKid());
        System.out.println("         • Status: " + originalMeta.getStatus());  // Should be ROTATED
        System.out.println("         • Next KID: " + originalMeta.getNextKid());  // Should point to successor

        System.out.println();
        System.out.println("[6/7] Retrieving successor key metadata...");

        ExportedKeySpec successorMeta = sdk.exportKey(successorKid);

        System.out.println("      Successor key:");
        System.out.println("         • KID: " + successorMeta.getKid());
        System.out.println("         • Status: " + successorMeta.getStatus());  // Should be ACTIVE
        System.out.println("         • Previous KID: " + successorMeta.getPreviousKid());  // Should point to original

        System.out.println();
        System.out.println("[7/7] Verifying bidirectional rotation chain...");

        if ("ROTATED".equals(originalMeta.getStatus()) &&
            successorKid.equals(originalMeta.getNextKid()) &&
            originalKid.equals(successorMeta.getPreviousKid())) {
            System.out.println("      ✅ Rotation chain verified:");
            System.out.println("         • Original → ROTATED status");
            System.out.println("         • Original.nextKid → Successor");
            System.out.println("         • Successor.previousKid → Original");
            System.out.println("         • Bidirectional linkage: VALID");
        } else {
            System.out.println("      ❌ Rotation chain verification failed");
        }
        System.out.println();
    }

    private static void demonstrateTransparentEncryption(AuthenticatedSdk sdk,
                                                         String originalKid,
                                                         String successorKid) throws Exception {

        System.out.println("[8/10] Creating test plaintext...");
        String plaintext = "Quantum-resistant data requiring key rotation";
        Path plaintextFile = TEMP_DIR.resolve("sc34_plaintext.txt");
        Path encryptedFile = TEMP_DIR.resolve("sc34_encrypted.jwe");

        FileIO.writeUtf8(plaintextFile, plaintext);
        System.out.println("       File: " + plaintextFile);
        System.out.println("       Size: " + plaintext.length() + " bytes");

        System.out.println();
        System.out.println("[9/10] Encrypting with ORIGINAL KID (should use successor)...");
        System.out.println("       Requested KID: " + originalKid + " (ROTATED)");

        EncryptResult result = sdk.encryptFileStream(originalKid, plaintextFile, encryptedFile);

        System.out.println("       ✅ Encryption successful");
        System.out.println("       Key requested: " + result.getKeyRequested());
        System.out.println("       Key actually used: " + result.getActualKeyUsed());
        System.out.println("       Algorithm: " + result.getAlgorithmUsed());

        System.out.println();
        System.out.println("[10/10] Analyzing transparent redirection...");

        if (originalKid.equals(result.getKeyRequested()) &&
            successorKid.equals(result.getActualKeyUsed())) {
            System.out.println("       ✅ Transparent redirection verified:");
            System.out.println("          • Client requested: " + originalKid);
            System.out.println("          • Server used: " + successorKid);
            System.out.println("          • Rotation transparent to client");
            System.out.println("          • No code changes required");
        } else {
            System.out.println("       ⚠️ Unexpected key usage");
        }

        System.out.println();
        System.out.println("       Security guarantee:");
        System.out.println("       • Encryption uses NEW rotated key");
        System.out.println("       • Historical data decrypts with OLD key");
        System.out.println("       • Zero-downtime rotation achieved");
        System.out.println();
    }
}

// Cannot change mode
Original: HYBRID_KEM_COMBINE
Successor: DUALSIGN
Result: InvalidInputException("Composite mode mismatch")

// Cannot change algorithms
Original: [X25519, ML-KEM-768]
Successor: [RSA-3072, ML-KEM-768]
Result: InvalidInputException("Component algorithm mismatch at index 0")

// Cannot downgrade to SIMPLE
Original: COMPOSITE (X25519 + ML-KEM-768)
Successor: SIMPLE (ML-KEM-1024)
Result: InvalidInputException("Cannot rotate COMPOSITE key to SIMPLE key")

// Flexible - can change algorithm family
sdk.createRotation("rsa-key",
    GenerateKeySpec.builder()
        .kty("ML-KEM")  // Different family OK
        .alg("ML-KEM-768")
        .build()
);

// Strict - must maintain mode and algorithms
sdk.createCompositeRotation("hybrid-key",
    GenerateCompositeKeySpec.builder()
        .mode(HYBRID_KEM_COMBINE)  // Must match
        .addComponent(classical(X25519))  // Must match
        .addComponent(pqc(ML_KEM_768))   // Must match
        .build()
);

ExportedKeySpec createCompositeRotation(
    String kid,
    GenerateCompositeKeySpec successor
) throws AnkaSecureSdkException