import { AesCtr, AesKw, computeJwkThumbprint, Hkdf, X25519 } from '@enbox/crypto';
import { Encoder } from './encoder.js';
import { KeyDerivationScheme } from './hd-key.js';
import { DwnError, DwnErrorCode } from '../core/dwn-error.js';
export var ContentEncryptionAlgorithm;
(function (ContentEncryptionAlgorithm) {
    ContentEncryptionAlgorithm["A256CTR"] = "A256CTR";
})(ContentEncryptionAlgorithm || (ContentEncryptionAlgorithm = {}));
export var KeyAgreementAlgorithm;
(function (KeyAgreementAlgorithm) {
    KeyAgreementAlgorithm["X25519HkdfSha256A256Kw"] = "X25519-HKDF-SHA256+A256KW";
})(KeyAgreementAlgorithm || (KeyAgreementAlgorithm = {}));
export const ROLE_AUDIENCE_DERIVATION_SCHEME = 'roleAudience';
export const SEAL_DERIVATION_SCHEME = 'seal';
const AES_256_KEY_LENGTH_BYTES = 32;
const AES_CTR_COUNTER_LENGTH_BYTES = 16;
const AES_CTR_COUNTER_LENGTH_BITS = 128;
const KEK_INFO_PREFIX = KeyAgreementAlgorithm.X25519HkdfSha256A256Kw;
function concatBytes(...arrays) {
    const byteLength = arrays.reduce((total, array) => total + array.byteLength, 0);
    const result = new Uint8Array(byteLength);
    let offset = 0;
    for (const array of arrays) {
        result.set(array, offset);
        offset += array.byteLength;
    }
    return result;
}
export class Encryption {
    static async encrypt(algorithm, keyBytes, initializationVector, plaintext) {
        Encryption.validateContentEncryptionParameters(algorithm, keyBytes, initializationVector);
        return AesCtr.encrypt({
            counter: initializationVector,
            data: plaintext,
            key: Encryption.toA256CtrJwk(keyBytes),
            length: AES_CTR_COUNTER_LENGTH_BITS,
        });
    }
    static async decrypt(algorithm, keyBytes, initializationVector, ciphertext) {
        Encryption.validateContentEncryptionParameters(algorithm, keyBytes, initializationVector);
        return AesCtr.decrypt({
            counter: initializationVector,
            data: ciphertext,
            key: Encryption.toA256CtrJwk(keyBytes),
            length: AES_CTR_COUNTER_LENGTH_BITS,
        });
    }
    static async encryptStream(algorithm, keyBytes, initializationVector, plaintextStream) {
        const plaintext = await Encryption.readStream(plaintextStream);
        const ciphertext = await Encryption.encrypt(algorithm, keyBytes, initializationVector, plaintext);
        return Encryption.toStream(ciphertext);
    }
    static async decryptStream(algorithm, keyBytes, initializationVector, ciphertextStream) {
        const ciphertext = await Encryption.readStream(ciphertextStream);
        const plaintext = await Encryption.decrypt(algorithm, keyBytes, initializationVector, ciphertext);
        return Encryption.toStream(plaintext);
    }
    static async wrapKey(input) {
        if (Encryption.isX25519KeyAgreementAlgorithm(input.keyInput.algorithm)) {
            return Encryption.wrapX25519Key(input);
        }
        throw new Error(Encryption.unsupportedKeyAgreementAlgorithmMessage(input.keyInput.algorithm));
    }
    static async wrapSeal(input) {
        if (!Encryption.isX25519KeyAgreementAlgorithm(input.keyInput.algorithm)) {
            throw new Error(Encryption.unsupportedKeyAgreementAlgorithmMessage(input.keyInput.algorithm));
        }
        Encryption.validateContentEncryptionKey(input.privateKeyBytes);
        const privateKey = input.ephemeralPrivateKey ?? await X25519.generateKey();
        const publicKey = await X25519.getPublicKey({ key: privateKey });
        const sharedSecret = await X25519.sharedSecret({
            privateKeyA: privateKey,
            publicKeyB: input.keyInput.publicKey,
        });
        const kek = await Encryption.deriveKek(sharedSecret, input.keyInput);
        const unwrappedKey = await AesKw.bytesToPrivateKey({ privateKeyBytes: input.privateKeyBytes });
        const encryptedKey = await AesKw.wrapKey({
            encryptionKey: Encryption.toA256KwJwk(kek),
            unwrappedKey,
        });
        return {
            algorithm: input.keyInput.algorithm,
            derivationScheme: SEAL_DERIVATION_SCHEME,
            encryptedKey: Encoder.bytesToBase64Url(encryptedKey),
            ephemeralPublicKey: publicKey,
            keyId: input.keyInput.keyId,
        };
    }
    static async wrapX25519Key(input) {
        Encryption.validateContentEncryptionKey(input.cek);
        const privateKey = input.ephemeralPrivateKey ?? await X25519.generateKey();
        const publicKey = await X25519.getPublicKey({ key: privateKey });
        const sharedSecret = await X25519.sharedSecret({
            privateKeyA: privateKey,
            publicKeyB: input.keyInput.publicKey,
        });
        const kek = await Encryption.deriveKek(sharedSecret, input.keyInput);
        const cekJwk = Encryption.toA256CtrJwk(input.cek);
        const kekJwk = { alg: 'A256KW', k: Encoder.bytesToBase64Url(kek), kty: 'oct' };
        const encryptedKey = await AesKw.wrapKey({ encryptionKey: kekJwk, unwrappedKey: cekJwk });
        return { encryptedKey, ephemeralPublicKey: publicKey };
    }
    static async unwrapKey(recipientPrivateKey, keyEncryption) {
        if (Encryption.isX25519KeyAgreementAlgorithm(keyEncryption.algorithm)) {
            return Encryption.unwrapX25519Key(recipientPrivateKey, keyEncryption);
        }
        throw new Error(Encryption.unsupportedKeyAgreementAlgorithmMessage(keyEncryption.algorithm));
    }
    static async unwrapSeal(input) {
        if (!Encryption.isX25519KeyAgreementAlgorithm(input.seal.algorithm)) {
            throw new Error(Encryption.unsupportedKeyAgreementAlgorithmMessage(input.seal.algorithm));
        }
        const keyEncryption = {
            ...input.seal,
            audienceKeyId: input.audienceKeyId,
            contextId: input.contextId,
            protocol: input.protocol,
            rolePath: input.rolePath,
        };
        const sharedSecret = await X25519.sharedSecret({
            privateKeyA: input.recipientPrivateKey,
            publicKeyB: input.seal.ephemeralPublicKey,
        });
        const kek = await Encryption.deriveKek(sharedSecret, keyEncryption);
        const unwrappedKey = await AesKw.unwrapKey({
            decryptionKey: Encryption.toA256KwJwk(kek),
            wrappedKeyAlgorithm: 'A256KW',
            wrappedKeyBytes: Encoder.base64UrlToBytes(input.seal.encryptedKey),
        });
        return AesKw.privateKeyToBytes({ privateKey: unwrappedKey });
    }
    static async unwrapX25519Key(recipientPrivateKey, keyEncryption) {
        const sharedSecret = await X25519.sharedSecret({
            privateKeyA: recipientPrivateKey,
            publicKeyB: keyEncryption.ephemeralPublicKey,
        });
        const kek = await Encryption.deriveKek(sharedSecret, keyEncryption);
        const kekJwk = { alg: 'A256KW', k: Encoder.bytesToBase64Url(kek), kty: 'oct' };
        const unwrappedJwk = await AesKw.unwrapKey({
            decryptionKey: kekJwk,
            wrappedKeyAlgorithm: ContentEncryptionAlgorithm.A256CTR,
            wrappedKeyBytes: Encoder.base64UrlToBytes(keyEncryption.encryptedKey),
        });
        return Encoder.base64UrlToBytes(unwrappedJwk.k);
    }
    static async buildEncryptionProperty(encryptionInput) {
        const algorithm = encryptionInput.algorithm ?? ContentEncryptionAlgorithm.A256CTR;
        Encryption.validateContentEncryptionParameters(algorithm, encryptionInput.key, encryptionInput.initializationVector);
        const keyEncryption = [];
        for (const keyInput of encryptionInput.keyEncryptionInputs) {
            keyEncryption.push(await Encryption.buildKeyEncryptionEntry(encryptionInput.key, keyInput));
        }
        return {
            algorithm,
            initializationVector: Encoder.bytesToBase64Url(encryptionInput.initializationVector),
            keyEncryption,
        };
    }
    static async getKeyId(publicKeyJwk) {
        return computeJwkThumbprint({ jwk: publicKeyJwk });
    }
    static validateEncryptionProperty(encryption) {
        const initializationVector = Encoder.base64UrlToBytes(encryption.initializationVector);
        if (initializationVector.byteLength !== AES_CTR_COUNTER_LENGTH_BYTES) {
            throw new DwnError(DwnErrorCode.RecordsWriteValidateIntegrityEncryptionInitializationVectorInvalid, `A256CTR initializationVector must decode to ${AES_CTR_COUNTER_LENGTH_BYTES} bytes.`);
        }
        for (const entry of encryption.keyEncryption) {
            Encryption.validateKeyEncryptionEntry(entry);
        }
    }
    static async buildKeyEncryptionEntry(cek, keyInput) {
        const { encryptedKey, ephemeralPublicKey } = await Encryption.wrapKey({
            cek,
            keyInput,
        });
        const common = {
            algorithm: keyInput.algorithm,
            encryptedKey: Encoder.bytesToBase64Url(encryptedKey),
            ephemeralPublicKey,
            keyId: keyInput.keyId,
        };
        if (keyInput.derivationScheme === KeyDerivationScheme.ProtocolPath) {
            return {
                ...common,
                derivationScheme: KeyDerivationScheme.ProtocolPath,
            };
        }
        return {
            ...common,
            derivationScheme: ROLE_AUDIENCE_DERIVATION_SCHEME,
            protocol: keyInput.protocol,
            rolePath: keyInput.rolePath,
        };
    }
    static validateKeyEncryptionEntry(entry) {
        if (Encryption.isX25519KeyAgreementAlgorithm(entry.algorithm)) {
            Encryption.validateX25519KeyEncryptionEntry(entry);
            return;
        }
        throw new DwnError(DwnErrorCode.RecordsWriteValidateIntegrityEncryptionEphemeralPublicKeyInvalid, Encryption.unsupportedKeyAgreementAlgorithmMessage(entry.algorithm));
    }
    static validateX25519KeyEncryptionEntry(entry) {
        if (entry.ephemeralPublicKey.kty !== 'OKP' || entry.ephemeralPublicKey.crv !== 'X25519') {
            throw new DwnError(DwnErrorCode.RecordsWriteValidateIntegrityEncryptionEphemeralPublicKeyInvalid, 'ephemeralPublicKey must be an OKP X25519 public key.');
        }
    }
    static async deriveKek(sharedSecret, keyEncryption) {
        Encryption.validateSharedSecret(sharedSecret);
        const info = Encryption.getKekInfo(keyEncryption);
        return Hkdf.deriveKeyBytes({
            baseKeyBytes: sharedSecret,
            hash: 'SHA-256',
            info,
            length: 256,
            salt: new Uint8Array(),
        });
    }
    static getKekInfo(keyEncryption) {
        if (keyEncryption.derivationScheme === SEAL_DERIVATION_SCHEME) {
            return JSON.stringify([
                KEK_INFO_PREFIX,
                SEAL_DERIVATION_SCHEME,
                keyEncryption.protocol,
                keyEncryption.rolePath,
                keyEncryption.contextId,
                keyEncryption.audienceKeyId,
            ]);
        }
        if (keyEncryption.derivationScheme === ROLE_AUDIENCE_DERIVATION_SCHEME) {
            return JSON.stringify([
                KEK_INFO_PREFIX,
                ROLE_AUDIENCE_DERIVATION_SCHEME,
                keyEncryption.protocol,
                keyEncryption.rolePath,
                keyEncryption.keyId,
            ]);
        }
        return JSON.stringify([KEK_INFO_PREFIX, KeyDerivationScheme.ProtocolPath, keyEncryption.keyId]);
    }
    static validateContentEncryptionParameters(algorithm, keyBytes, initializationVector) {
        if (algorithm !== ContentEncryptionAlgorithm.A256CTR) {
            throw new Error(`Unsupported content encryption algorithm: ${algorithm}`);
        }
        Encryption.validateContentEncryptionKey(keyBytes);
        if (initializationVector.byteLength !== AES_CTR_COUNTER_LENGTH_BYTES) {
            throw new Error(`A256CTR initializationVector must be ${AES_CTR_COUNTER_LENGTH_BYTES} bytes.`);
        }
    }
    static validateContentEncryptionKey(keyBytes) {
        if (keyBytes.byteLength !== AES_256_KEY_LENGTH_BYTES) {
            throw new Error(`A256CTR content encryption key must be ${AES_256_KEY_LENGTH_BYTES} bytes.`);
        }
    }
    static validateSharedSecret(sharedSecret) {
        if (sharedSecret.every((byte) => byte === 0)) {
            throw new Error('X25519 shared secret MUST NOT be all zeros.');
        }
    }
    static isX25519KeyAgreementAlgorithm(algorithm) {
        return algorithm === KeyAgreementAlgorithm.X25519HkdfSha256A256Kw;
    }
    static unsupportedKeyAgreementAlgorithmMessage(algorithm) {
        return `Unsupported key agreement algorithm: ${String(algorithm)}`;
    }
    static toA256CtrJwk(keyBytes) {
        return { alg: ContentEncryptionAlgorithm.A256CTR, k: Encoder.bytesToBase64Url(keyBytes), kty: 'oct' };
    }
    static toA256KwJwk(keyBytes) {
        return { alg: 'A256KW', k: Encoder.bytesToBase64Url(keyBytes), kty: 'oct' };
    }
    static async readStream(stream) {
        const reader = stream.getReader();
        const chunks = [];
        for (;;) {
            const { done, value } = await reader.read();
            if (done) {
                break;
            }
            chunks.push(value);
        }
        return concatBytes(...chunks);
    }
    static toStream(data) {
        return new ReadableStream({
            start(controller) {
                controller.enqueue(data);
                controller.close();
            }
        });
    }
}
//# sourceMappingURL=encryption.js.map