/*
 * Copyright (c)2013-2020 ZeroTier, Inc.
 *
 * Use of this software is governed by the Business Source License included
 * in the LICENSE.TXT file in the project's root directory.
 *
 * Change Date: 2026-01-01
 *
 * On the date above, in accordance with the Business Source License, use
 * of this software will be governed by version 2.0 of the Apache License.
 */
/****/

#ifndef ZT_N_CAMOPATTERN_HPP
#define ZT_N_CAMOPATTERN_HPP

#include <stdint.h>
#include <array>
#include <unordered_map>
#include <mutex>
#include <random>

#include "Address.hpp"
#include "Buffer.hpp"
#include "RuntimeEnvironment.hpp"

#define BYTES_IN_WORD (sizeof(uint32_t) / sizeof(uint8_t))
#define PATTERN_COUNT 16
#define NO_CAMO PATTERN_COUNT

#define ZT_CAMO_DEFAULT_WORDSTR "DLGE"

/**
 * Camo functions debug trace macro
*/
#ifdef CAMO_TRACE
    #define CT(...) do { \
        printf("%s:%d %s: ", __FILE__, __LINE__, __PRETTY_FUNCTION__); \
        printf(__VA_ARGS__); \
        printf("\n"); \
    } while(0)
#else
    #define CT(...) ((void)0)
#endif

/**
 * Type shorthands
*/
namespace ZeroTier {

enum class CamoLevel;
typedef std::array<uint8_t, BYTES_IN_WORD> CamoPatternBytes;
typedef std::array<CamoPatternBytes, PATTERN_COUNT> CamoPatternArray;
typedef std::unordered_map<CamoPatternBytes, size_t> CamoIndexMap;
typedef std::unordered_map<Address, CamoLevel> KnownHostsMap;

}

/**
 * Hash functions for the respective unordered_maps
*/
namespace std {
    template<>
    struct hash<ZeroTier::CamoPatternBytes> {
        size_t operator()(const ZeroTier::CamoPatternBytes& bytes) const {
            uint32_t word = 0;
            for (const auto& byte : bytes) {
                word <<= 8;
                word |= byte;
            }
            return std::hash<uint32_t>{}(word);
        }
    };

    template<>
    struct hash<ZeroTier::Address> {
        size_t operator()(const ZeroTier::Address& address) const {
            return std::hash<uint64_t>{}(address.toInt());
        }
    };
}

namespace ZeroTier {

/**
 * Camouflage level enum
 */
enum class CamoLevel {
    NONE,
    NODE,
    CONTROLLER,
    MOON,
    PLANET,
    INAPPLICABLE
};

enum class CamoRelayRule {
    LEAVE,
    KNOWNHOSTS,
    STRIP,
    APPLY
};

/**
 * Helper class for packet camouflage operations
 */
class CamoPattern
{
    /**
     * Check if buffer has camouflage applied
     *
     * @param buffer Buffer to check
     * @return True if buffer has camouflage
     */
    template<unsigned int C>
    static size_t getCamoIndex(const Buffer<C> &buffer)
    {
        size_t result = NO_CAMO;
        CamoPatternBytes camo;
        if (buffer.size() > BYTES_IN_WORD * 2)
        {
            for (size_t i = 0; i < BYTES_IN_WORD; i++)
            {
                camo[i] = buffer[i] ^ buffer[i + BYTES_IN_WORD];
            }
            auto it = camoIndices.find(camo);
            if (it != camoIndices.end())
            {
                result = it->second;
                CT("CAMO DETECTED, INDEX: %u", result);
            }
            else
            {
                CT("CAMO NOT DETECTED");
            }
        }
        return result;
    }

    /**
     * Check the host camo level
     *
     * @param host Destination address
     * @param RR RuntimeEnvironment pointer
     * @return Camo Level for this host
     */
    static CamoLevel getCamoLevel(const Address host, const RuntimeEnvironment * const RR);

public:
    /**
     *
     * @param host Destination address
     * @param RR RuntimeEnvironment pointer
     * @return True if host requires camo
     */
    static bool isCamoRequired(const Address host, const RuntimeEnvironment * const RR, const bool hadCamo = false, const bool isRelay = false);

    /**
     * Apply camouflage to buffer
     *
     * This increases buffer length by adding ID word to the end
     *
     * @param buffer Buffer to apply camouflage to
     */
    template<unsigned int C>
    static void applyCamo(Buffer<C> &buffer)
    {
        size_t camoIndex = getCamoIndex(buffer);
        if (isInitialized && (camoIndex == NO_CAMO)) { // Only apply if not already applied
            CT("PACKET CONTENTS BEFORE APPLYING CAMO:");
            buffer.dump();
            camoIndex = std::minstd_rand(std::time(nullptr))() % PATTERN_COUNT;
            CamoPatternBytes camo = camoValues[camoIndex];

            // Increase buffer size first to avoid overflow
            size_t originalSize = buffer.size();
            buffer.setSize(originalSize + BYTES_IN_WORD);
            uint8_t * const data = reinterpret_cast<uint8_t *>(buffer.unsafeData());

            // Copy the second word to the end
            for (size_t i = 0; i < BYTES_IN_WORD; i++) {
                data[i + originalSize] = data[i + BYTES_IN_WORD];
            }

            // Copy the first word on the second word's place
            for (size_t i = 0; i < BYTES_IN_WORD; i++) {
                data[i + BYTES_IN_WORD] = data[i];
            }

            // Apply XOR to the rest of the buffer
            for (size_t i = BYTES_IN_WORD; i < buffer.size(); i++) {
                data[i] ^= camo[i % BYTES_IN_WORD];
            }

            CT("PACKET CONTENTS AFTER APPLYING CAMO:");
            buffer.dump();
        }
    }

    /**
     * Remove camouflage from buffer
     *
     * This decreases buffer length by removing stored ID word from the end
     *
     * @param buffer Buffer to remove camouflage from
     */
    template<unsigned int C>
    static bool stripCamo(Buffer<C> &buffer)
    {
        bool result = false;
        size_t camoIndex = NO_CAMO;
        if (isInitialized)
        {
            camoIndex = getCamoIndex(buffer);
        }
        if (camoIndex != NO_CAMO) {
            CT("PACKET CONTENTS BEFORE STRIPPING CAMO:");
            buffer.dump();
            uint8_t * const data = reinterpret_cast<uint8_t *>(buffer.unsafeData());
            CamoPatternBytes camo = camoValues[camoIndex];
            for (size_t i = BYTES_IN_WORD * 2; i < buffer.size(); i++)
            {
              data[i] ^= camo[i % BYTES_IN_WORD];
            }
            size_t storedWordIndex = buffer.size() - BYTES_IN_WORD;
            for (size_t i = 0; i < BYTES_IN_WORD; i++)
            {
                data[i + BYTES_IN_WORD] = data[i + storedWordIndex];
            }
            buffer.setSize(buffer.size() - BYTES_IN_WORD);
            result = true;
            CT("PACKET CONTENTS AFTER STRIPPING CAMO:");
            buffer.dump();
        }
        return result;
    }

    static void init(CamoLevel level, uint32_t word, KnownHostsMap hosts, CamoRelayRule rule);


private:
    static bool isInitialized;
    static CamoLevel camoLevel;
    static uint32_t camoWord;
    static CamoRelayRule relayRule;
    static CamoPatternArray camoValues;
    static CamoIndexMap camoIndices;
    static std::mutex camoMutex;
    static KnownHostsMap knownHosts;
};

} // namespace ZeroTier

// Implementation of Buffer::dump() method
template<unsigned int C>
void ZeroTier::Buffer<C>::dump() const
{
#ifdef CAMO_TRACE
    const unsigned char *bytes = reinterpret_cast<const unsigned char *>(data());
    const unsigned int size = this->size();

    for (unsigned int i = 0; i < size; i++) {
        printf("%02x", bytes[i]);

        if ((i + 1) % 16 == 0) {
            printf("\n");
        } else if (i < size - 1) {
            printf(" ");
        }
    }

    // Add newline if last line wasn't complete
    if (size % 16 != 0) {
        printf("\n");
    }
#endif
}

#endif // ZT_CAMOPATTERN_HPP