berstream.hpp

#pragma once

#include "dangerous/snmp/logger.hpp"

#include "bytestream.hpp"

#include <vector>

namespace dangerous { namespace snmp {

class BerStream {
public:
    BerStream( ByteStream* byteStream );

    bool copyTo( std::vector<char>& buffer );

    unsigned int remainingReadLength() const { return _byteStream->remainingReadLength(); }
    
    unsigned int remainingWriteLength() const { return _byteStream->remainingWriteLength(); }

    /*
     * Peek functions.
     */

    bool peekType( uint8_t& type );
    
    /*
     * Read functions.
     * All of these functions will advance the read head.
     */

    bool readType( uint8_t& type );

    bool readLength( unsigned int& length );

    template<typename EncodingClass>
    bool read( typename EncodingClass::value_type& value );

    bool readBytes( char* buffer, unsigned int bufferSize, unsigned int& bytesRead );

    bool readFullTlv( std::vector<char>& buffer );

    /*
     * Write functions.
     * All of these functions will advance the write head.
     */

    bool writeType( uint8_t type );

    bool writeLength( unsigned int encodedLength );

    template<typename EncodingClass>
    bool write( const typename EncodingClass::value_type& value );

    bool writeBytes( const char* buffer, unsigned int bufferSize, unsigned int& bytesWritten );

protected:
    ByteStream* _byteStream;
};

template<typename EncodingClass>
bool BerStream::read( typename EncodingClass::value_type& value ) {
    uint8_t type = 0;
    bool status = readType( type );
    if( ! status ) {
        if( logger.system( Logger::BYTESTREAM ) ) {
            logger.out() << "BerStream: read<" << EncodingClass::NAME << ">: Could not read the type." << std::endl;
        }
        return false;
    }

    if( logger.system( Logger::BYTESTREAM ) ) {
        logger.out() << "BerStream: read<" << EncodingClass::NAME << ">: Looking for a type of " << EncodingClass::TYPE << "." << std::endl;
    }
    if( type != EncodingClass::TYPE ) {
        return false;
    }

    unsigned int increment = 0;
    status = readLength( increment );
    if( ! status ) {
        if( logger.system( Logger::BYTESTREAM ) ) {
            logger.out() << "BerStream: read<" << EncodingClass::NAME << ">: Could not read the length." << std::endl;
        }
        return false;
    }
    if( logger.system( Logger::BYTESTREAM ) ) {
        logger.out() << "BerStream: read<" << EncodingClass::NAME << ">: Got a length of " << increment << "." << std::endl;
    }
    if( increment < 1 ) {
        return true;
    }

    if( logger.system( Logger::BYTESTREAM ) ) {
        logger.out() << "BerStream: read<" << EncodingClass::NAME << ">: Reading " << increment << " bytes." << std::endl;
    }
    unsigned int bytesRead = 0;
    char buffer[ increment ];
    status = readBytes( buffer, increment, bytesRead );
    if( ! status ) {
        if( logger.system( Logger::BYTESTREAM ) ) {
            logger.out() << "BerStream: read<" << EncodingClass::NAME << ">: Could not read the value bytes." << std::endl;
        }
        return false;
    }
    status = EncodingClass::read( value, buffer, increment );
    if( ! status ) {
        if( logger.system( Logger::BYTESTREAM ) ) {
            logger.out() << "BerStream: read<" << EncodingClass::NAME << ">: Could not parse the value bytes." << std::endl;
        }
        return false;
    }

    return true;
}

template<typename EncodingClass>
bool BerStream::write( const typename EncodingClass::value_type& value ) {
    if( logger.system( Logger::BYTESTREAM ) ) {
        logger.out() << "BerStream: write<" << EncodingClass::NAME << ">: Writing a type of " << EncodingClass::TYPE << "." << std::endl;
    }
    bool status = writeType( EncodingClass::TYPE );
    if( ! status ) {
        if( logger.system( Logger::BYTESTREAM ) ) {
            logger.out() << "BerStream: write<" << EncodingClass::NAME << ">: Could not write the type." << std::endl;
        }
        return false;
    }

    unsigned int encodedLength = EncodingClass::length( value );
    if( logger.system( Logger::BYTESTREAM ) ) {
        logger.out() << "BerStream: write<" << EncodingClass::NAME << ">: Got a length of " << encodedLength << "." << std::endl;
    }
    status = writeLength( encodedLength );
    if( ! status ) {
        if( logger.system( Logger::BYTESTREAM ) ) {
            logger.out() << "BerStream: write<" << EncodingClass::NAME << ">: Could not write the length." << std::endl;
        }
        return false;
    }

    if( logger.system( Logger::BYTESTREAM ) ) {
        logger.out() << "BerStream: write<" << EncodingClass::NAME << ">: Writing " << encodedLength << " bytes." << std::endl;
    }
    
    char buffer[ encodedLength ];
    status = EncodingClass::write( value, buffer, encodedLength );
    if( ! status ) {
        if( logger.system( Logger::BYTESTREAM ) ) {
            logger.out() << "BerStream: write<" << EncodingClass::NAME << ">: Could not encode the value." << std::endl;
        }
        return false;
    }
    
    unsigned int bytesWritten = 0;
    status = writeBytes( buffer, encodedLength, bytesWritten );
    if( ! status ) {
        if( logger.system( Logger::BYTESTREAM ) ) {
            logger.out() << "BerStream: write<" << EncodingClass::NAME << ">: Could not write the value." << std::endl;
        }
        return false;
    }
    
    return true;
}

} }