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.AnkaSecureSdk;
import co.ankatech.ankasecure.sdk.exception.AnkaSecureSdkException;
import co.ankatech.ankasecure.sdk.model.*;
import co.ankatech.ankasecure.sdk.util.FileIO;

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)
 * Demonstrates quantum-resistant key rotation with validation
 */
public final class ExampleScenario34 {

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

    public static void main(String[] args) {
        System.out.println("===== SCENARIO 34: COMPOSITE KEY ROTATION =====");

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

            runScenario(sdk);

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

    private static void runScenario(AnkaSecureSdk sdk) throws Exception {

        // PHASE 1: Generate Original Composite Key
        String originalKid = generateOriginalCompositeKey(sdk);

        // PHASE 2: Rotate to Successor
        String successorKid = rotateToSuccessor(sdk, originalKid);

        // PHASE 3: Verify Rotation Chain
        verifyRotationChain(sdk, originalKid, successorKid);

        // PHASE 4: Demonstrate Transparent Encryption
        demonstrateTransparentEncryption(sdk, originalKid, successorKid);
    }

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

        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);

        ExportedKeySpec result = sdk.generateCompositeKey(spec);

        System.out.println("✅ Original composite key generated");
        System.out.println("   KID: " + kid);
        System.out.println("   Algorithm: " + result.getAlg());
        System.out.println("   Status: " + result.getStatus());

        return kid;
    }

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

        // Successor must maintain same mode and algorithms
        GenerateCompositeKeySpec successorSpec = new GenerateCompositeKeySpec()
                .setKid(successorKid)
                .setMode(GenerateCompositeKeySpec.Mode.HYBRID_KEM_COMBINE)  // MUST match
                .addComponent(ComponentSpec.classical(ClassicalCompositeAlgorithm.X25519))  // MUST match
                .addComponent(ComponentSpec.pqc(PqcCompositeAlgorithm.ML_KEM_768))  // MUST match
                .setKdf(Kdf.HKDF_SHA256)
                .setMaxUsageLimit(2000000);  // Can change

        System.out.println("Rotating composite key...");
        ExportedKeySpec successor = sdk.createCompositeRotation(originalKid, successorSpec);

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

        return successorKid;
    }

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

        // Get metadata for both keys
        ExportedKeySpec originalMeta = sdk.exportKey(originalKid);
        ExportedKeySpec successorMeta = sdk.exportKey(successorKid);

        System.out.println("✅ Rotation chain verified:");
        System.out.println("   Original: " + originalKid);
        System.out.println("      • Status: " + originalMeta.getStatus());
        System.out.println("      • Next KID: " + originalMeta.getNextKid());
        System.out.println("   Successor: " + successorKid);
        System.out.println("      • Status: " + successorMeta.getStatus());
        System.out.println("      • Previous KID: " + successorMeta.getPreviousKid());
    }

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

        String plaintext = "Data requiring quantum-resistant key rotation";
        Path plaintextFile = TEMP_DIR.resolve("sc34_plaintext.txt");
        Path encryptedFile = TEMP_DIR.resolve("sc34_encrypted.jwe");

        FileIO.writeUtf8(plaintextFile, plaintext);

        // Encrypt with original KID (server redirects to successor)
        EncryptResult result = sdk.encryptFileStream(originalKid, plaintextFile, encryptedFile);

        System.out.println("✅ Transparent encryption:");
        System.out.println("   Requested: " + originalKid);
        System.out.println("   Actually used: " + result.getActualKeyUsed());
        System.out.println("   Zero-downtime rotation achieved");
    }
}

// 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