// SPDX-License-Identifier: GPL-2.0-only
/*
 * Digital Voice Modem - Common Library
 * MIT Open Source. Use is subject to license terms.
 * DO NOT ALTER OR REMOVE COPYRIGHT NOTICES OR THIS FILE HEADER.
 *
 *  Copyright (C) 2025 Bryan Biedenkapp, N2PLL
 *
 */
/**
 * @defgroup p25_crypto Project 25 Cryptography
 * @brief Defines and implements cryptography routines for P25.
 * @ingroup p25
 * 
 * @file Crypto.h
 * @ingroup p25
 * @file Crypto.cpp
 * @ingroup p25
 */
#if !defined(__P25_CRYPTO_H__)
#define __P25_CRYPTO_H__

#include "common/Defines.h"
#include "common/p25/P25Defines.h"
#include "common/Utils.h"

#include <random>

namespace p25
{
    namespace crypto
    {
        // ---------------------------------------------------------------------------
        //  Class Declaration
        // ---------------------------------------------------------------------------

        /**
         * @brief Project 25 Cryptography.
         * @ingroup p25_crypto
         */
        class HOST_SW_API P25Crypto {
        public:
            /**
             * @brief Initializes a new instance of the P25Crypto class.
             */
            P25Crypto();
            /**
             * @brief Finalizes a instance of the P25Crypto class.
             */
            ~P25Crypto();

            /**
             * @brief Helper given to generate a new initial seed MI.
             * @param mi 
             */
            void generateMI();
            /**
             * @brief Helper given the last MI, generate the next MI using LFSR.
             */
            void generateNextMI();

            /**
             * @brief Helper to check if there is a valid encryption keystream.
             * @return bool True, if there is a valid keystream, otherwise false.
             */
            bool hasValidKeystream();
            /**
             * @brief Helper to generate the encryption keystream.
             */
            void generateKeystream();
            /**
             * @brief Helper to reset the encryption keystream.
             */
            void resetKeystream();

            /**
             * @brief Helper to crypt a P25 TEK with the given AES-256 KEK.
             * @note This assumes AES-256 KEK and will not work with other key types.
             * @param kek Key Encryption Key
             * @param tek Traffic Encryption Key
             * @param tekLen Traffic Encryption Key Length
             * @returns UInt8Array Buffer containing the encryption wrapped TEK.
             */
            UInt8Array cryptAES_TEK(const uint8_t* kek, uint8_t* tek, uint8_t tekLen);
            /**
             * @brief Helper to decrypt a P25 TEK with the given AES-256 KEK.
             * @note This assumes AES-256 KEK and will not work with other key types.
             * @param kek Key Encryption Key
             * @param tek Wrapped Traffic Encryption Key
             * @param tekLen Wrapped Traffic Encryption Key Length
             * @returns UInt8Array Buffer containing the decrypted TEK.
             */
            UInt8Array decryptAES_TEK(const uint8_t* kek, uint8_t* tek, uint8_t tekLen);
            /**
             * @brief Helper to generate a P25 KMM CBC MAC key with the given AES-256 KEK.
             * @note This assumes AES-256 KEK and will not work with other key types.
             * @param kek Key Encryption Key
             * @param msg KMM Message
             * @param msgLen Message Length
             * @returns UInt8Array Buffer containing the derived MAC key.
             */
            UInt8Array cryptAES_KMM_CBC_KDF(const uint8_t* kek, const uint8_t* msg, uint16_t msgLen);
            /**
             * @brief Helper to generate a P25 KMM CBC-MAC with the given AES-256 CBC-MAC key.
             * @note This assumes AES-256 KEK and will not work with other key types.
             * @param macKey CBC-MAC Derived Key
             * @param msg KMM Message
             * @param msgLen Message Length
             * @returns UInt8Array Buffer containing the derived MAC key.
             */
            UInt8Array cryptAES_KMM_CBC(const uint8_t* macKey, const uint8_t* msg, uint16_t msgLen);

            /**
             * @brief Helper to generate a P25 KMM CMAC MAC key with the given AES-256 KEK.
             * @note This assumes AES-256 KEK and will not work with other key types.
             * @param kek Key Encryption Key
             * @param msg KMM Message
             * @param msgLen Message Length
             * @param hasMN Message has a message number.
             * @returns UInt8Array Buffer containing the derived MAC key.
             */
            UInt8Array cryptAES_KMM_CMAC_KDF(const uint8_t* kek, const uint8_t* msg, uint16_t msgLen, bool hasMN);
            /**
             * @brief Helper to generate a P25 KMM CMAC with the given AES-256 CMAC key.
             * @note This assumes AES-256 KEK and will not work with other key types.
             * @param macKey CMAC Derived Key
             * @param msg KMM Message
             * @param msgLen Message Length
             * @returns UInt8Array Buffer containing the derived MAC key.
             */
            UInt8Array cryptAES_KMM_CMAC(const uint8_t* macKey, const uint8_t* msg, uint16_t msgLen);

            /**
             * @brief Helper to crypt a P25 PDU frame using AES-256.
             * @param frame Data frame
             * @param frameLen Data frame Length
             */
            void cryptAES_PDU(uint8_t* frame, uint8_t frameLen);

            /**
             * @brief Helper to crypt P25 IMBE audio using DES.
             * @param imbe Buffer containing IMBE to crypt.
             * @param duid P25 DUID.
             */
            void cryptDES_IMBE(uint8_t* imbe, P25DEF::DUID::E duid);
            /**
             * @brief Helper to crypt P25 IMBE audio using AES-256.
             * @param imbe Buffer containing IMBE to crypt.
             * @param duid P25 DUID.
             */
            void cryptAES_IMBE(uint8_t* imbe, P25DEF::DUID::E duid);
            /**
             * @brief Helper to crypt P25 IMBE audio using ARC4.
             * @param imbe Buffer containing IMBE to crypt.
             * @param duid P25 DUID.
             */
            void cryptARC4_IMBE(uint8_t* imbe, P25DEF::DUID::E duid);

            /**
             * @brief Helper to check if there is a valid encryption message indicator.
             * @return bool True, if there is a valid encryption message indicator, otherwise false.
             */
            bool hasValidMI();
            /**
             * @brief Sets the encryption message indicator.
             * @param[in] mi Buffer containing the 9-byte Message Indicator.
             */
            void setMI(const uint8_t* mi);
            /**
             * @brief Gets the encryption message indicator.
             * @param[out] mi Buffer containing the 9-byte Message Indicator.
             */
            void getMI(uint8_t* mi) const;
            /**
             * @brief Clears the encryption message indicator.
             */
            void clearMI();

            /**
             * @brief Sets the encryption key.
             * @param[in] mi Buffer containing the encryption key.
             * @param[in] len Length of the key.
             */
            void setKey(const uint8_t* key, uint8_t len);
            /**
             * @brief Gets the encryption key,
             * @param[out] mi Buffer containing the encryption key.
             */
            void getKey(uint8_t* key) const;
            /**
             * @brief Clears the stored encryption key.
             */
            void clearKey();

        public:
            /**
             * @brief Traffic Encryption Key Algorithm ID.
             */
            DECLARE_PROPERTY(uint8_t, tekAlgoId, TEKAlgoId);
            /**
             * @brief Traffic Encryption Key ID.
             */
            DECLARE_PROPERTY(uint16_t, tekKeyId, TEKKeyId);

            /**
             * @brief Traffic Encryption Key Length.
             */
            DECLARE_RO_PROPERTY(uint8_t, tekLength, TEKLength);

        private:
            uint8_t* m_keystream;
            uint32_t m_keystreamPos;
    
            uint8_t* m_mi;

            UInt8Array m_tek;

            std::mt19937 m_random;

            /**
             * @brief Helper to step the linear feedback shift register (LFSR).
             * @note This uses the polynomial: x^64 + x^62 + x^46 + x^38 + x^27 + x^15 + 1
             * @param lfsr Linear feedback shift register value.
             * @return uint64_t The next LFSR value.
             */
            uint64_t stepLFSR(uint64_t& lfsr);
            /**
             * @brief Expands the 9-byte MI into a proper 16-byte IV.
             * @return uint8_t* Buffer containing expanded 16-byte IV.
             */
            uint8_t* expandMIToIV();
        };
    } // namespace crypto
} // namespace p25

#endif // __P25_CRYPTO_H__
