Flow 4 – Streaming Re-encrypt RSA-3072 ➜ ML-KEM-768

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

package co.ankatech.ankasecure.sdk.examples;

import co.ankatech.ankasecure.sdk.AuthenticatedSdk;
import co.ankatech.ankasecure.sdk.model.DecryptResultMetadata;
import co.ankatech.ankasecure.sdk.model.EncryptResult;
import co.ankatech.ankasecure.sdk.model.ReencryptResult;
import co.ankatech.ankasecure.sdk.util.FileIO;

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

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

/**
 * Scenario 4 — Streaming Re-encrypt (RSA-3072 → ML-KEM-768).
 *
 * <p>This scenario demonstrates how to migrate classical RSA-encrypted data to
 * a post-quantum ML-KEM-encrypted format <em>without exposing plaintext</em>.
 * All cryptographic operations use the <strong>streaming</strong> API, meaning
 * the service exchanges a two-part <em>detached&nbsp;JWE</em>
 * (General-JSON header + raw ciphertext stream).</p>
 *
 * <p>The SDK is data-plane-only: it does not create keys. This scenario operates
 * on the <strong>pre-provisioned granted cross-kid pair</strong> of the
 * demo playground &mdash; source {@code cli-reencrypt2-source}
 * ({@code RSA-3072}) and target {@code cli-reencrypt2-target} ({@code ML-KEM-768}).
 * The playground actor holds the cross-kid REENCRYPT grant on this pair
 * plus the atomic encrypt/decrypt grants on each endpoint.</p>
 *
 * <ol>
 *   <li>Resolve the granted {@code RSA-3072} source and {@code ML-KEM-768} target kids.</li>
 *   <li>Stream-encrypt a file under the source (detached&nbsp;JWE).</li>
 *   <li>Re-encrypt the ciphertext on the server (RSA&nbsp;&rarr;&nbsp;KEM).</li>
 *   <li>Stream-decrypt with the KEM target key.</li>
 *   <li>Validate round-trip integrity.</li>
 * </ol>
 *
 * <p>Temporary artefacts are placed under <kbd>temp_files/</kbd>.</p>
 *
 * <p><b>Implementation notes (Java&nbsp;21+):</b></p>
 * <ul>
 *   <li>All file operations rely on the {@link java.nio.file.Path} API.</li>
 *   <li>Encoding is explicitly set to UTF-8 for deterministic behaviour.</li>
 *   <li>Directory creation delegates to {@link ExampleUtil#ensureTempDir(Path)}.</li>
 * </ul>
 *
 * <p><b>Thread-safety:</b> this class is stateless and immutable.</p>
 *
 * @see ExampleUtil
 * @see ExamplePlaygroundKeys
 * @see AuthenticatedSdk
 */
public final class ExampleScenario4 {

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

    /**
     * Runs the RSA-3072 to ML-KEM-768 streaming re-encryption 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(final String[] args) {

        System.out.println("===== SCENARIO 4 START =====");
        System.out.println("""
                Purpose :
                  * RSA-3072 -> ML-KEM-768 streaming re-encryption
                  * Uses the pre-provisioned granted cross-kid pair
                    (cli-reencrypt2-source / cli-reencrypt2-target)
                  * Demonstrates detached-JWE pipeline (encrypt -> re-encrypt -> decrypt)
                Steps   :
                  1) Resolve granted RSA-3072 source + ML-KEM-768 target kids
                  2) Create dummy payload
                  3) Encrypt payload under source (detached JWE)
                  4) Re-encrypt ciphertext stream to target
                  5) Decrypt KEM ciphertext under target
                  6) Validate integrity
                --------------------------------------------------------------""");

        try {
            prepareWorkingDir();

            Properties            props      = loadProperties();
            AuthenticatedSdk      sdk        = authenticate(props);
            ExamplePlaygroundKeys playground = ExamplePlaygroundKeys.from(props, sdk);

            runScenario(sdk, playground);

        } catch (Exception ex) {
            fatal("Scenario 4 failed", ex);
        }

        System.out.println("===== SCENARIO 4 END =====");
    }

    /**
     * Executes the streaming re-encryption scenario on the granted cross-kid pair.
     *
     * <ol>
     *   <li>Create large plaintext file (1 MB).</li>
     *   <li>Resolve the granted RSA-3072 source and ML-KEM-768 target kids.</li>
     *   <li>Stream-encrypt under the RSA-3072 source.</li>
     *   <li>Re-encrypt RSA ciphertext to ML-KEM (server-side).</li>
     *   <li>Stream-decrypt under the ML-KEM-768 target.</li>
     *   <li>Validate binary integrity via stream comparison.</li>
     * </ol>
     *
     * @param sdk        an authenticated {@link AuthenticatedSdk} instance;
     *                   must not be {@code null}
     * @param playground the pre-provisioned playground key resolver; must not be {@code null}
     * @throws Exception on any I/O or cryptographic failure
     */
    private static void runScenario(final AuthenticatedSdk sdk,
                                    final ExamplePlaygroundKeys playground) throws Exception {

        /* 1 -- prepare large plaintext file (e.g., 1 Mb) --------------------- */
        Path plainFile = TEMP_DIR.resolve("scenario4_plain_large.bin");
        byte[] dummy = new byte[1 * 1024 * 1024]; // zero-filled 1 Mb buffer
        FileIO.writeBytes(plainFile, dummy);
        System.out.println("[1] Plaintext ready          -> " + plainFile.toAbsolutePath());

        /* 2 -- resolve granted cross-kid pair (data-plane SDK: no key creation) */
        final String rsaKid = playground.requireKid("cli-reencrypt2-source"); // RSA-3072
        final String kemKid = playground.requireKid("cli-reencrypt2-target"); // ML-KEM-768
        System.out.println("[2] Source kid (RSA-3072)    -> " + rsaKid);
        System.out.println("[2] Target kid (ML-KEM-768)  -> " + kemKid);

        /* 3 -- streaming encrypt under RSA-3072 source ----------------------- */
        Path rsaCipher = TEMP_DIR.resolve("scenario4_rsa_large.enc");
        EncryptResult rsaEncMeta = sdk.encryptFileStream(rsaKid, plainFile, rsaCipher);
        System.out.println("[3] RSA ciphertext written   -> " + rsaCipher.toAbsolutePath());
        printEncryptMeta(rsaEncMeta);

        /* 4 -- one-pass streaming re-encrypt RSA -> ML-KEM ------------------- */
        Path kemCipher = TEMP_DIR.resolve("scenario4_kem_large.enc");
        ReencryptResult reMeta = sdk.reencryptFileStream(kemKid, null, rsaCipher, kemCipher);
        System.out.println("[4] Re-encrypted file        -> " + kemCipher.toAbsolutePath());
        printReencryptMeta(reMeta);

        /* 5 -- streaming decrypt under ML-KEM-768 target -------------------- */
        Path decFile = TEMP_DIR.resolve("scenario4_large_decrypted.bin");
        DecryptResultMetadata decMeta = sdk.decryptFileStream(kemCipher, decFile);
        System.out.println("[5] Decrypted file          -> " + decFile.toAbsolutePath());
        printDecryptMeta(decMeta);

        /* 6 -- validation via streaming compare ------------------------------ */
        try (InputStream in1 = Files.newInputStream(plainFile);
             InputStream in2 = Files.newInputStream(decFile)) {

            boolean match = streamsAreEqual(in1, in2);
            System.out.println(match
                    ? "[6] Validation OK - plaintext matches."
                    : "[6] WARNING - plaintext mismatch!");
        }
    }
}