/**
 * Rawr-X3DH -- eXtended 3-way Diffie-Hellman
 *
 * Specification by Open Whisper Systems <https://signal.org/docs/specifications/x3dh/>
 * Powered by Libsodium <https://libsodium.gitbook.io/doc/>
 *
 * Implemented by Soatok Dreamseeker <https://soatok.blog>
 *
 * ................................:.................
 * .............................-+yd-................
 * ............/+:-.....+/oys++://:m:................
 * --........../y///oyssyyyyhddh+-:y/................
 * --------.....o--+syyso/syyyhho:--+..........:-....
 * ----------.....:/ssss+ooyoosyo//yo.--------oy:....
 * --------------:+//+//++:`-/o-syyy/-------+yo------
 * --------------:oy++:s. ++:+: `ys/------/ss:-------
 * ---------------:+++syh--/++++oss//---:oy+---------
 * ----------------:syyhhysosssyyyyhso/+yo:----------
 * ----------------::shddyyhyyyyshdhyyyy/------------
 * ----------------:shhhyyyssssssoyhhhyhho/::--------
 * ::::::---------:+shhhddd+o+++++yddhhhhhyyyso+::::-
 * ::::::::::::o+oyssyhyhdh+///:+hhoshhhhhhhhyo+:::::
 * ::::::::::::+syyssss/:yyoo+sydo/o+s+/+osyysoo/::::
 * ::::::::::::/+ssyyyyy/:oyyhhhs/ss/y/::::::::::::::
 * :::::::::::::::/+syhhhsyhhhhyyss+oo/::::::::::::::
 * :::::::::/o+/:::::/+syhddddysyso+so/:::::////:::::
 * ::::::+yhhyhhs::::::/yhdddddssso/ss:::::::///:::::
 * :::::::::hhhh+/+/:/shoshyysss/  `+s:::::://::::/::
 * ::::::::+hhhho:+sshhsyhhhs+:.     `-//::::::::::::
 * ::::::::ohhyoo+oyhyyhhhyyssoo/:-`    .:/::::::::::
 * :::::::::syso+syhhhhhhhhhhhhhhyyyo:`   ./:::::::::
 * :::::::::--://+o++osyo+yhhhhhhhhhhyys/`  :::::::::
 * ::::::--.:/+/::::-::::yhhhhhhhhhhhyyy+.  :::::::::
 * :------://:::::----:::/yhyyyyyyyyyys+`   :+:::::::
 * ------::-----------:shyhhyyyyss+/:-...-::+//::::::
 * ------------------/yhhhhhhyyyssso+::::::::::::::::
 * -----------------+yyyhhhhhhhyyssso+/---------:::::
 * ---------------/syyys/yhhhhhhyyyysss+-------------
 * .............:syyyyo---oyhhhhhhhhyyyhs------------
 * ...........-oyyyyyo....-+syyhyyhhhhddy------------
 * ...........syyyyys-......-::::+:////:.------------
 * ...........yyyyys:............-...............----
 * ...........+sss:..................................
 * .....````````.``..................................
 *
 */
import type {
    Ed25519SecretKey, X25519SecretKey
} from 'sodium-plus'
import {
    CryptographyKey,
    Ed25519PublicKey,
    SodiumPlus,
    X25519PublicKey
} from 'sodium-plus'
import type {
    KeyDerivationFunction,
    SymmetricEncryptionInterface
} from './src/symmetric'
import {
    blakeKdf,
    SymmetricCrypto
} from './src/symmetric'
import type {
    SessionKeyManagerInterface,
    IdentityKeyManagerInterface
} from './src/persistence'
import {
    DefaultSessionKeyManager,
    DefaultIdentityKeyManager
} from './src/persistence'
import {
    concat,
    generateKeyPair,
    generateBundle,
    signBundle,
    verifyBundle,
    wipe
} from './src/util'

/**
 * Initial server info.
 *
 * Contains the information necessary to complete
 * the X3DH handshake from a sender's side.
 */
export type InitServerInfo = {
    IdentityKey:string,
    SignedPreKey:{
        Signature:string,
        PreKey:string
    },
    OneTimeKey?:string
};

/**
 * Initial information about a sender
 */
export type InitSenderInfo = {
    Sender:string,
    IdentityKey:string,
    EphemeralKey:string,
    OneTimeKey?:string,
    CipherText:string
};

/**
 * Send a network request to the server to obtain the public keys needed
 * to complete the sender's handshake.
 */
export type InitClientFunction = (id:string)=>Promise<InitServerInfo>;

/**
 * Signed key bundle.
 */
export type SignedBundle = { signature:string, bundle:string[] };

/**
 * Initialization information for receiving a handshake message.
 */
type RecipientInitWithSK = {
    IK:Ed25519PublicKey,
    EK:X25519PublicKey,
    SK:CryptographyKey,
    OTK?:string
};

/**
 * Pluggable X3DH implementation, powered by libsodium.
 */
export class X3DH {
    encryptor:SymmetricEncryptionInterface
    kdf:KeyDerivationFunction
    identityKeyManager:IdentityKeyManagerInterface
    sessionKeyManager:SessionKeyManagerInterface
    sodium?:SodiumPlus

    constructor (
        identityKeyManager?:IdentityKeyManagerInterface,
        sessionKeyManager?:SessionKeyManagerInterface,
        encryptor?:SymmetricEncryptionInterface,
        kdf?:KeyDerivationFunction
    ) {
        if (!sessionKeyManager) {
            sessionKeyManager = new DefaultSessionKeyManager()
        }
        if (!identityKeyManager) {
            identityKeyManager = new DefaultIdentityKeyManager()
        }
        if (!encryptor) {
            encryptor = new SymmetricCrypto()
        }
        if (!kdf) {
            kdf = blakeKdf
        }
        this.encryptor = encryptor
        this.kdf = kdf
        this.sessionKeyManager = sessionKeyManager
        this.identityKeyManager = identityKeyManager
    }

    /**
     * @returns {SodiumPlus}
     */
    async getSodium ():Promise<SodiumPlus> {
        if (!this.sodium) {
            this.sodium = await SodiumPlus.auto()
        }
        return this.sodium
    }

    /**
     * Generates and signs a bundle of one-time keys.
     *
     * Useful for pushing more OTKs to the server.
     *
     * @param {Ed25519SecretKey} signingKey
     * @param {number} numKeys
     */
    async generateOneTimeKeys (
        signingKey:Ed25519SecretKey,
        numKeys:number = 100
    ):Promise<SignedBundle> {
        const sodium = await this.getSodium()
        const bundle = await generateBundle(numKeys)
        const publicKeys = bundle.map(x => x.publicKey)
        const signature = await signBundle(signingKey, publicKeys)
        await this.identityKeyManager.persistOneTimeKeys(bundle)

        // Hex-encode all the public keys
        const encodedBundle : string[] = []
        for (const pk of publicKeys) {
            encodedBundle.push(await sodium.sodium_bin2hex(pk.getBuffer()))
        }

        return {
            signature: await sodium.sodium_bin2hex(signature),
            bundle: encodedBundle
        }
    }

    /**
     * Get the shared key when sending an initial message.
     *
     * @param {InitServerInfo} res
     * @param {Ed25519SecretKey} senderKey
     */
    async initSenderGetSK (
        res:InitServerInfo,
        senderKey:Ed25519SecretKey
    ):Promise<RecipientInitWithSK> {
        const sodium = await this.getSodium()
        const identityKey = new Ed25519PublicKey(
            await sodium.sodium_hex2bin(res.IdentityKey)
        )
        const signedPreKey = new X25519PublicKey(
            await sodium.sodium_hex2bin(res.SignedPreKey.PreKey)
        )
        const signature = await sodium.sodium_hex2bin(res.SignedPreKey.Signature)

        // Check signature
        const valid = await verifyBundle(identityKey, [signedPreKey], signature)
        if (!valid) {
            throw new Error('Invalid signature')
        }
        const ephemeral = await generateKeyPair()
        const ephSecret = ephemeral.secretKey
        const ephPublic = ephemeral.publicKey

        // Turn the Ed25519 keys into X25519 keys for X3DH:
        const senderX = await sodium.crypto_sign_ed25519_sk_to_curve25519(senderKey)
        const recipientX = await sodium.crypto_sign_ed25519_pk_to_curve25519(identityKey)

        // See the X3DH specification to really understand this part:
        const DH1 = await sodium.crypto_scalarmult(senderX, signedPreKey)
        const DH2 = await sodium.crypto_scalarmult(ephSecret, recipientX)
        const DH3 = await sodium.crypto_scalarmult(ephSecret, signedPreKey)
        let SK
        if (res.OneTimeKey) {
            const DH4 = await sodium.crypto_scalarmult(
                ephSecret,
                new X25519PublicKey(await sodium.sodium_hex2bin(res.OneTimeKey))
            )
            SK = new CryptographyKey(
                Buffer.from(await this.kdf(
                    concat(
                        DH1.getBuffer(),
                        DH2.getBuffer(),
                        DH3.getBuffer(),
                        DH4.getBuffer()
                    )
                ))
            )
            await wipe(DH4)
        } else {
            SK = new CryptographyKey(
                Buffer.from(await this.kdf(
                    concat(
                        DH1.getBuffer(),
                        DH2.getBuffer(),
                        DH3.getBuffer()
                    )
                ))
            )
        }

        // Wipe DH keys since we have SK
        await wipe(DH1)
        await wipe(DH2)
        await wipe(DH3)
        await wipe(ephSecret)
        await wipe(senderX)

        return {
            IK: identityKey,
            EK: ephPublic,
            SK,
            OTK: res.OneTimeKey
        }
    }

    /**
     * Initialize for sending.
     *
     * @param {string} recipientIdentity
     * @param {InitClientFunction} getServerResponse
     * @param {string|Buffer} message
     */
    async initSend (
        recipientIdentity:string,
        getServerResponse:InitClientFunction,
        message:string|Buffer
    ):Promise<InitSenderInfo> {
        const sodium = await this.getSodium()

        // Get the identity key for the sender:
        const senderIdentity = await this.identityKeyManager.getMyIdentityString()
        const identity = await this.identityKeyManager.getIdentityKeypair()
        const senderSecretKey = identity.identitySecret
        const senderPublicKey = identity.identityPublic

        // Stub out a call to get the server response:
        const response = await getServerResponse(recipientIdentity)

        // Get the shared symmetric key (and other handshake data):
        const { IK, EK, SK, OTK } = await this.initSenderGetSK(response, senderSecretKey)

        // Get the assocData for AEAD:
        const assocData = await sodium.sodium_bin2hex(
            Buffer.concat([senderPublicKey.getBuffer(), IK.getBuffer()])
        )

        // Set the session key (as a sender):
        await this.sessionKeyManager.setSessionKey(recipientIdentity, SK, false)
        await this.sessionKeyManager.setAssocData(recipientIdentity, assocData)
        return {
            Sender: senderIdentity,
            IdentityKey: await sodium.sodium_bin2hex(senderPublicKey.getBuffer()),
            EphemeralKey: await sodium.sodium_bin2hex(EK.getBuffer()),
            OneTimeKey: OTK,
            CipherText: await this.encryptor.encrypt(
                message,
                await this.sessionKeyManager.getEncryptionKey(recipientIdentity),
                assocData
            )
        }
    }

    /**
     * Get the shared key when receiving an initial message.
     *
     * @param {InitSenderInfo} req
     * @param {Ed25519SecretKey} identitySecret
     * @param preKeySecret
     */
    async initRecvGetSk (
        req:InitSenderInfo,
        identitySecret:Ed25519SecretKey,
        preKeySecret:X25519SecretKey
    ) {
        const sodium = await this.getSodium()

        // Decode strings
        const senderIdentityKey = new Ed25519PublicKey(
            await sodium.sodium_hex2bin(req.IdentityKey),
        )
        const ephemeral = new X25519PublicKey(
            await sodium.sodium_hex2bin(req.EphemeralKey),
        )

        // Ed25519 -> X25519
        const senderX = await sodium.crypto_sign_ed25519_pk_to_curve25519(senderIdentityKey)
        const recipientX = await sodium.crypto_sign_ed25519_sk_to_curve25519(identitySecret)

        // See the X3DH specification to really understand this part:
        const DH1 = await sodium.crypto_scalarmult(preKeySecret, senderX)
        const DH2 = await sodium.crypto_scalarmult(recipientX, ephemeral)
        const DH3 = await sodium.crypto_scalarmult(preKeySecret, ephemeral)

        let SK
        if (req.OneTimeKey) {
            const DH4 = await sodium.crypto_scalarmult(
                await this.identityKeyManager.fetchAndWipeOneTimeSecretKey(req.OneTimeKey),
                ephemeral
            )
            SK = new CryptographyKey(
                Buffer.from(await this.kdf(
                    concat(
                        DH1.getBuffer(),
                        DH2.getBuffer(),
                        DH3.getBuffer(),
                        DH4.getBuffer()
                    )
                ))
            )
            await wipe(DH4)
        } else {
            SK = new CryptographyKey(
                Buffer.from(await this.kdf(
                    concat(
                        DH1.getBuffer(),
                        DH2.getBuffer(),
                        DH3.getBuffer()
                    )
                ))
            )
        }
        // Wipe DH keys since we have SK
        await wipe(DH1)
        await wipe(DH2)
        await wipe(DH3)
        await wipe(recipientX)
        return {
            Sender: req.Sender,
            SK,
            IK: senderIdentityKey
        }
    }

    /**
     * Initialize keys for receiving an initial message.
     * Returns the initial plaintext message on success.
     * Throws on failure.
     *
     * @param {InitSenderInfo} req
     * @returns {(string|Buffer)[]}
     */
    async initRecv (req:InitSenderInfo):Promise<(string|Buffer)[]> {
        const sodium = await this.getSodium()
        const { identitySecret, identityPublic } = await this.identityKeyManager.getIdentityKeypair()
        const { preKeySecret } = await this.identityKeyManager.getPreKeypair()
        const { Sender, SK, IK } = await this.initRecvGetSk(
            req,
            identitySecret,
            preKeySecret
        )
        const assocData = await sodium.sodium_bin2hex(
            Buffer.from(concat(IK.getBuffer(), identityPublic.getBuffer()))
        )
        try {
            await this.sessionKeyManager.setSessionKey(Sender, SK, true)
            await this.sessionKeyManager.setAssocData(Sender, assocData)
            return [
                Sender,
                await this.encryptor.decrypt(
                    req.CipherText,
                    await this.sessionKeyManager.getEncryptionKey(Sender, true),
                    assocData
                )
            ]
        } catch (e) {
            // Decryption failure! Destroy the session.
            await this.sessionKeyManager.destroySessionKey(Sender)
            throw e
        }
    }

    /**
     * Encrypt the next message to send to the recipient.
     *
     * @param {string} recipient
     * @param {string|Buffer} message
     * @returns {string}
     */
    async encryptNext (recipient:string, message:string|Buffer):Promise<string> {
        return this.encryptor.encrypt(
            message,
            await this.sessionKeyManager.getEncryptionKey(recipient, false),
            await this.sessionKeyManager.getAssocData(recipient)
        )
    }

    /**
     * Decrypt the next message received by the sender.
     *
     * @param {string} sender
     * @param {string} encrypted
     * @returns {string|Buffer}
     */
    async decryptNext (sender:string, encrypted:string):Promise<string|Buffer> {
        return this.encryptor.decrypt(
            encrypted,
            await this.sessionKeyManager.getEncryptionKey(sender, true),
            await this.sessionKeyManager.getAssocData(sender)
        )
    }

    /**
     * Sets the identity string for the current user.
     *
     * @param {string} id
     */
    async setIdentityString (id:string):Promise<void> {
        return this.identityKeyManager.setMyIdentityString(id)
    }
}

/* Let's make sure we export the interfaces/etc we use. */
export * from './src/symmetric'
export * from './src/persistence'
export * from './src/util'