apps/elftosb/common/AESCounter.h

/*
 * File:    AESCounter.h
 *
 * Copyright (c) Freescale Semiconductor, Inc. All rights reserved.
 * See included license file for license details.
 */
#if !defined(_AESCounter_h_)
#define _AESCounter_h_

#include <string.h>
#include <iostream>
#include "Random.h"

typedef uint8_t aes128_counter_t[16];

class AESCounterBase
{
public:
    void _readFromStream(std::istream &stream, unsigned bytes, uint8_t *buffer);

    void _writeToStream(std::ostream &stream, unsigned bytes, const uint8_t *buffer) const;
};

template <int S>
class AESCounter : public AESCounterBase
{
public:
    typedef uint8_t counter_t[S / 8];

public:
    AESCounter() { memset(m_counter, 0, sizeof(m_counter)); }
    AESCounter(const counter_t &counter) { memcpy(m_counter, &counter, sizeof(m_counter)); }
    // \brief Constructor taking a counter value pointer.
    AESCounter(const counter_t *counter) { memcpy(m_counter, counter, sizeof(m_counter)); }
    AESCounter(std::istream &stream) { readFromStream(stream); }
    AESCounter(const AESCounter<S> &other) { memcpy(m_counter, other.m_counter, sizeof(m_counter)); }
    ~AESCounter() { memset(m_counter, 0, sizeof(m_counter)); }
    void randomize()
    {
        RandomNumberGenerator rng;
        rng.generateBlock(m_counter, sizeof(m_counter));
    }

    void readFromStream(std::istream &stream)
    {
        _readFromStream(stream, S / 8, reinterpret_cast<uint8_t *>(&m_counter));
    }

    void writeToStream(std::ostream &stream) const
    {
        _writeToStream(stream, S / 8, reinterpret_cast<const uint8_t *>(&m_counter));
    }

    void incrementCounter(unsigned val)
    {
        for (unsigned j = 0; j < val; j++)
        {
            for (unsigned i = sizeof(AESCounter<S>::counter_t) - 1, carry = 1; (i >= 0) && carry; i--)
            {
                carry = !++m_counter[i];
            }
        }
    }


    inline const counter_t &getCounter() const { return m_counter; }
    inline void getCounter(counter_t *counter) const { memcpy(counter, m_counter, sizeof(m_counter)); }
    inline void setCounter(const counter_t &counter) { memcpy(m_counter, &counter, sizeof(m_counter)); }
    inline void setCounter(const counter_t *counter) { memcpy(m_counter, counter, sizeof(m_counter)); }
    inline void setCounter(const AESCounter<S> &counter) { memcpy(m_counter, counter.m_counter, sizeof(m_counter)); }


    const AESCounter<S> &operator=(const AESCounter<S> &counter)
    {
        setCounter(counter);
        return *this;
    }
    const AESCounter<S> &operator=(const counter_t &counter)
    {
        setCounter(counter);
        return *this;
    }
    const AESCounter<S> &operator=(const counter_t *counter)
    {
        setCounter(counter);
        return *this;
    }

    operator const counter_t &() const { return m_counter; }
    operator const counter_t *() const { return m_counter; }
    friend std::ostream &operator<<(std::ostream &os, const AESCounter<S> &counter)
    {
        counter.writeToStream(os);
        return os;
    }

protected:
    counter_t m_counter; 
};

typedef AESCounter<128> AES128Counter;

#endif // _AESCounter_h_