Digital Watermarking (C++, Qt3, gnulib)

/**
 * Copyright (C) 2009 Kamil Dudka <xdudka00@stud.fit.vutbr.cz>
 *
 * This file is part of kry2.
 *
 * kry2 is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * any later version.
 *
 * kry2 is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with kry2.  If not, see <http://www.gnu.org/licenses/>.
 */
 
#include <cassert>
#include <iomanip>
#include <iostream>
#include <stdexcept>
#include <stdint.h>
#include <vector>
 
#include <qcolor.h>
#include <qimage.h>
 
extern "C" {
#   include "md5.h"
#   include "hamming.h"
}
 
#include "Color.h"
 
#ifndef DEBUG_HAM
#   define DEBUG_HAM 0
#endif
 
#ifndef DEBUG_HASH
#   define DEBUG_HASH 0
#endif
 
#ifndef DEBUG_LEN
#   define DEBUG_LEN 0
#endif
 
/// byte array
typedef std::vector<unsigned char> TData;
 
/**
 * Generic interface for block based I/O.
 */
class IBlkIO {
    public:
        virtual ~IBlkIO() { }
 
        /**
         * Return size of one block. All blocks ought to be the same size.
         */
        virtual size_t blkSize() const = 0;
 
        /**
         * Return total count of blocks available to read/write.
         */
        virtual size_t blkCnt() const = 0;
 
        /**
         * Read one block.
         * @param blk Index of block to read. It has to be in range
         * <0, blkCnt()-1>
         * @param data @b Already @b sized byte array to store result to.
         * The array has to be the size of blkSize().
         */
        virtual void read(unsigned blk, TData &data) = 0;
 
        /**
         * Write one block.
         * @param blk Index of block to write to. It has to be in range
         * <0, blkCnt()-1>
         * @param data Array of data to write.
         * The array has to be the size of blkSize().
         */
        virtual void write(unsigned blk, const TData &data) = 0;
};
 
/**
 * IBlkIO implementation reading/writing from/to image.
 */
class ImageBlkIO: public IBlkIO {
    public:
        /**
         * @param image Instance of QImage used to read/write data from/to.
         * The QImage object has to be valid till ImageBlkIO destruction.
         */
        ImageBlkIO(QImage &image);
 
        // see IBlkIO dox
        virtual ~ImageBlkIO();
        virtual size_t blkSize() const;
        virtual size_t blkCnt() const;
        virtual void read(unsigned blk, TData &);
        virtual void write(unsigned blk, const TData &);
 
    private:
        struct Private;
        Private *d;
};
 
// this class is not intended to be exported
class HamCodec {
    public:
        HamCodec():
            hasError_(false)
        {
        }
 
        unsigned char encode(unsigned char);
        unsigned char decode(unsigned char);
 
        bool hasError() const { return hasError_; }
        void clearError() { hasError_ = 0; }
 
    private:
        bool hasError_;
};
 
/**
 * IBlkIO decorator reading/writing ECC (Hamming) code.
 * It decreases count of available blocks two times.
 */
class HamDecorator: public IBlkIO {
    public:
        HamDecorator(IBlkIO &parent);
        virtual ~HamDecorator();
        virtual size_t blkSize() const;
        virtual size_t blkCnt() const;
        virtual void read(unsigned blk, TData &);
        virtual void write(unsigned blk, const TData &);
 
    public:
        bool hasError() const;
        void clearError();
 
    private:
        struct Private;
        Private *d;
};
 
/**
 * IBlkIO decorator reading/writing inputLength field.
 * It decrements count of available blocks by one.
 */
class LenDecorator: public IBlkIO {
    public:
        LenDecorator(IBlkIO &parent);
        virtual ~LenDecorator();
        virtual size_t blkSize() const;
        virtual size_t blkCnt() const;
        virtual void read(unsigned blk, TData &);
        virtual void write(unsigned blk, const TData &);
 
    public:
        /**
         * Return inputLength field stored in data.
         */
        size_t getLength();
 
        /**
         * Store given inputLength field to data.
         * @param len Value ought to be stored.
         */
        void setLength(size_t len);
 
    private:
        struct Private;
        Private *d;
};
 
/**
 * IBlkIO decorator reading/writing data hash using md5 hash function.
 * It decreases count of available blocks by max 16/blkSize() + 1.
 */
class HashDecorator: public IBlkIO {
    public:
        HashDecorator(IBlkIO &parent);
        virtual ~HashDecorator();
        virtual size_t blkSize() const;
        virtual size_t blkCnt() const;
        virtual void read(unsigned blk, TData &);
        virtual void write(unsigned blk, const TData &);
 
    public:
        /**
         * Hash length [in bytes].
         */
        static const size_t HASH_LEN_IN_BYTES = 16;
 
        /**
         * Compute hash for stored data.
         * @param buff @b Already @b sized byte array to store result to.
         */
        void computeHash(TData &buff);
 
        /**
         * Write computed hash to data.
         */
        void wrtHash();
 
        /**
         * Compare stored hash with the computed one.
         * @return Return true if hashes are equal.
         */
        bool chkHash();
 
    private:
        struct Private;
        Private *d;
};
 
// /////////////////////////////////////////////////////////////////////////////
// ImageBlkIO implementation
struct ImageBlkIO::Private {
    QImage          &image;
    const unsigned  width;
    const unsigned  blkSize;
    const unsigned  blkCnt;
 
    Private(QImage &image_):
        image(image_),
        width(image.width()),
        blkSize(3),
        blkCnt(width * image.height() / 8)
    {
    }
};
 
ImageBlkIO::ImageBlkIO(QImage &image):
    d(new Private (image))
{
}
 
ImageBlkIO::~ImageBlkIO() {
    delete d;
}
 
size_t ImageBlkIO::blkSize() const {
    return d->blkSize;
}
 
size_t ImageBlkIO::blkCnt() const {
    return d->blkCnt;
}
 
void ImageBlkIO::read(unsigned blk, TData &data) {
    assert(blk < d->blkCnt);
    assert(data.size() == d->blkSize);
 
    uint32_t serialized = 0;
 
    // for each of 8 pixels
    const unsigned offset = blk * 8;
    for (int i = 7; 0 <= i; --i) {
        const int pixel = offset + i;
        const int x = pixel % d->width;
        const int y = pixel / d->width;
 
        // read colors
        int r, g, b;
        QColor color(d->image.pixel(x, y));
        color.getRgb(&r, &g, &b);
 
        // shift LSB from each color to bit-stream
        serialized <<= 1; serialized |= b & 1;
        serialized <<= 1; serialized |= g & 1;
        serialized <<= 1; serialized |= r & 1;
    }
 
    // split bit-stream to bytes and save to byte array
    static const uint32_t UCHAR_MASK = (1 << 8) - 1;
    data[2] = UCHAR_MASK & serialized; serialized >>= 8;
    data[1] = UCHAR_MASK & serialized; serialized >>= 8;
    data[0] = UCHAR_MASK & serialized;
}
 
void ImageBlkIO::write(unsigned blk, const TData &data) {
    assert(blk < d->blkCnt);
    assert(data.size() == d->blkSize);
 
    // join byte array to bit-stream
    static const uint32_t UCHAR_MASK = (1 << 8) - 1;
    uint32_t serialized = UCHAR_MASK & data[0];
    serialized <<= 8;
    serialized |= UCHAR_MASK & data[1];
    serialized <<= 8;
    serialized |= UCHAR_MASK & data[2];
 
    // form each of 8 pixels
    const unsigned offset = blk * 8;
    for (unsigned i = 0; i < 8; ++i) {
        const int pixel = offset + i;
        const int x = pixel % d->width;
        const int y = pixel / d->width;
 
        // read colors
        int r, g, b;
        QColor color(d->image.pixel(x, y));
        color.getRgb(&r, &g, &b);
 
        // shift bit-stream to LSB of each color
        r &= -2; r |= serialized & 1; serialized >>= 1;
        g &= -2; g |= serialized & 1; serialized >>= 1;
        b &= -2; b |= serialized & 1; serialized >>= 1;
 
        // save colors
        QColor newColor(r, g, b);
        d->image.setPixel(x, y, newColor.rgb());
    }
}
 
// /////////////////////////////////////////////////////////////////////////////
// HamCodec implementation
unsigned char HamCodec::encode(unsigned char in) {
    return HammingMatrixEncode(in);
}
 
unsigned char HamCodec::decode(unsigned char in) {
    unsigned char out = HammingMatrixDecode(in);
    if (in != HammingMatrixEncode(out)) {
        hasError_ = true;
#if DEBUG_HAM
        std::cerr << "in1 = 0x" << std::hex << std::setw(2)
                << static_cast<unsigned int>(in) << std::endl;
        std::cerr << "out = 0x" << std::hex << std::setw(2)
                << static_cast<unsigned int>(out) << std::endl;
        std::cerr << "in2 = 0x" << std::hex << std::setw(2)
                << static_cast<unsigned int>(HammingMatrixEncode(out))
                << std::endl << std::endl;
#endif
    }
    return out;
}
 
// /////////////////////////////////////////////////////////////////////////////
// HamDecorator implementation
struct HamDecorator::Private {
    IBlkIO          &parent;
    const unsigned  blkSize;
    const unsigned  blkCnt;
    HamCodec        codec;
 
    Private(IBlkIO &parent_):
        parent(parent_),
        blkSize(parent_.blkSize()),
        blkCnt(parent_.blkCnt() >> 1)
    {
    }
};
 
HamDecorator::HamDecorator(IBlkIO &parent):
    d(new Private (parent))
{
}
 
HamDecorator::~HamDecorator() {
    delete d;
}
 
size_t HamDecorator::blkSize() const {
    return d->blkSize;
}
 
size_t HamDecorator::blkCnt() const {
    return d->blkCnt;
}
 
void HamDecorator::read(unsigned blk, TData &data) {
    assert(blk < d->blkCnt);
    assert(data.size() == d->blkSize);
 
    TData encoded1(d->blkSize);
    TData encoded2(d->blkSize);
 
    // read two blocks from input
    const unsigned offset = blk << 1;
    d->parent.read(offset, encoded1);
    d->parent.read(offset + 1, encoded2);
 
    // for each byte of output
    for (unsigned i = 0; i < data.size(); ++i) {
        unsigned char &current = data[i];
 
        // decode two bytes and join into one output byte
        current = d->codec.decode(encoded1[i]);
        current |= d->codec.decode(encoded2[i]) << 4;
    }
}
 
void HamDecorator::write(unsigned blk, const TData &data) {
    assert(blk < d->blkCnt);
    assert(data.size() == d->blkSize);
 
    TData encoded1(d->blkSize);
    TData encoded2(d->blkSize);
 
    // for each byte of input
    for (unsigned i = 0; i < data.size(); ++i) {
        unsigned char current = data[i];
 
        // encode both nibbles and save as two bytes
        encoded1[i] = d->codec.encode(current & 0x0F);
        current >>= 4;
        encoded2[i] = d->codec.encode(current & 0x0F);
    }
 
    // save the two encoded blocks
    const unsigned offset = blk << 1;
    d->parent.write(offset, encoded1);
    d->parent.write(offset + 1, encoded2);
}
 
bool HamDecorator::hasError() const {
    return d->codec.hasError();
}
 
void HamDecorator::clearError() {
    d->codec.clearError();
}
 
// /////////////////////////////////////////////////////////////////////////////
// LenDecorator implementation
struct LenDecorator::Private {
    IBlkIO          &parent;
    const unsigned  lenBlks;
    const unsigned  blkCnt;
 
    Private(IBlkIO &parent_):
        parent(parent_),
        // FIXME: support unlimited length of input
        lenBlks(1),
        blkCnt(parent_.blkCnt() - lenBlks)
    {
    }
};
 
LenDecorator::LenDecorator(IBlkIO &parent):
    d(new Private (parent))
{
}
 
LenDecorator::~LenDecorator() {
    delete d;
}
 
size_t LenDecorator::blkSize() const {
    return d->parent.blkSize();
}
 
size_t LenDecorator::blkCnt() const {
    return d->blkCnt;
}
 
void LenDecorator::read(unsigned blk, TData &data) {
    assert(blk < d->blkCnt);
    assert(data.size() == d->parent.blkSize());
    d->parent.read(d->lenBlks + blk, data);
}
 
void LenDecorator::write(unsigned blk, const TData &data) {
    assert(blk < d->blkCnt);
    assert(data.size() == d->parent.blkSize());
    d->parent.write(d->lenBlks + blk, data);
}
 
size_t LenDecorator::getLength() {
    TData chunk(d->parent.blkSize());
    d->parent.read(0, chunk);
 
#if DEBUG_LEN
    std::cerr << "--- read length " << std::flush;
#endif
    // FIXME: support unlimited length of input
    size_t len = 0;
    for(int j = std::min(chunk.size(), sizeof(len)) - 1; 0 <= j; --j) {
        len <<= 8;
        len |= chunk[j];
#if DEBUG_LEN
        std::cerr << ".";
#endif
    }
#if DEBUG_LEN
    std::cerr << " " << len << std::endl;
#endif
    return len;
}
 
void LenDecorator::setLength(size_t len) {
    static const uint32_t UCHAR_MASK = (1 << 8) - 1;
    TData chunk(d->parent.blkSize());
 
#if DEBUG_LEN
    std::cerr << "--- writing length: " << len << std::flush;
#endif
    // FIXME: support unlimited length of input
    for(unsigned j = 0; j < chunk.size(); ++j) {
        chunk[j] = (j < sizeof(len))
            ? (len & UCHAR_MASK)
            : 0;
        len >>= 8;
#if DEBUG_LEN
        std::cerr << ".";
#endif
    }
    d->parent.write(0, chunk);
#if DEBUG_LEN
    std::cerr << std::endl;
#endif
}
 
 
// /////////////////////////////////////////////////////////////////////////////
// HashDecorator implementation
struct HashDecorator::Private {
    IBlkIO          &parent;
    const unsigned  blkSize;
    const unsigned  hashBlks;
    const unsigned  blkCnt;
 
    Private(IBlkIO &parent_):
        parent(parent_),
        blkSize(parent_.blkSize()),
        hashBlks((HASH_LEN_IN_BYTES % blkSize)
                ? (HASH_LEN_IN_BYTES / blkSize + 1)
                : (HASH_LEN_IN_BYTES / blkSize)),
        blkCnt(parent_.blkCnt() - hashBlks)
    {
    }
};
 
HashDecorator::HashDecorator(IBlkIO &parent):
    d(new Private (parent))
{
}
 
HashDecorator::~HashDecorator() {
    delete d;
}
 
size_t HashDecorator::blkSize() const {
    return d->blkSize;
}
 
size_t HashDecorator::blkCnt() const {
    return d->blkCnt;
}
 
void HashDecorator::read(unsigned blk, TData &data) {
    assert(blk < d->blkCnt);
    assert(data.size() == d->blkSize);
    d->parent.read(d->hashBlks + blk, data);
}
 
void HashDecorator::write(unsigned blk, const TData &data) {
    assert(blk < d->blkCnt);
    assert(data.size() == d->blkSize);
    d->parent.write(d->hashBlks + blk, data);
}
 
void HashDecorator::computeHash(TData &buff) {
    assert(buff.size() == HASH_LEN_IN_BYTES);
 
    // initialize crypto/md5 context
    md5_ctx ctx;
    md5_init_ctx(&ctx);
 
    // working space
    TData data(d->blkSize);
    unsigned char *raw_data = new unsigned char[d->blkSize];
 
    // for each block reachable from outside
    for (unsigned i = 0; i < d->blkCnt; i++) {
        this->read(i, data);
 
        // copy byte array from STL vector to native C++ array
        for (unsigned j = 0; j < d->blkSize; ++j)
            raw_data[j] = data[j];
 
        // process block by crypto/md5
        md5_process_bytes(raw_data, d->blkSize, &ctx);
    }
    delete []raw_data;
 
    // obtain crypto/md5 result
    unsigned char *result = new unsigned char[HASH_LEN_IN_BYTES];
    md5_finish_ctx(&ctx, result);
 
    // copy native C++ array to STL vector
    for (unsigned i = 0; i < HASH_LEN_IN_BYTES; ++i)
        buff[i] = result[i];
    delete []result;
}
 
void HashDecorator::wrtHash() {
    TData buff(HASH_LEN_IN_BYTES);
    this->computeHash(buff);
 
    TData chunk(d->parent.blkSize());
 
#if DEBUG_HASH
    std::cerr << "---  writing hash: 0x";
    for (unsigned i = 0; i < HashDecorator::HASH_LEN_IN_BYTES; ++i)
        std::cerr << std::hex << std::setw(2)
            << static_cast<unsigned int>(buff[i]);
    std::cerr << std::endl;
#endif
 
    // for each block
    unsigned pos = 0;
    for(unsigned i = 0; i < d->hashBlks; ++i) {
 
        // for each byte of block
        for(unsigned j = 0; j < chunk.size(); ++j) {
            chunk[j] = (pos < HASH_LEN_IN_BYTES)
                ? buff[pos++]
                : 0;
        }
        d->parent.write(i, chunk);
    }
}
 
bool HashDecorator::chkHash() {
    // read stored hash
    TData hash(HASH_LEN_IN_BYTES);
    TData chunk(d->parent.blkSize());
    unsigned pos = 0;
 
    // for each block
    for(unsigned i = 0; i < d->hashBlks; ++i) {
        d->parent.read(i, chunk);
 
        // for each byte of block
        for(unsigned j = 0; j < chunk.size(); ++j) {
            hash[pos] = chunk[j];
            if (HASH_LEN_IN_BYTES <= ++pos)
                break;
        }
    }
#if DEBUG_HASH
    std::cerr << "---     read hash: 0x";
    for (unsigned i = 0; i < HashDecorator::HASH_LEN_IN_BYTES; ++i)
        std::cerr << std::hex << std::setw(2)
            << static_cast<unsigned int>(hash[i]);
    std::cerr << std::endl;
#endif
 
    // compute current hash
    TData buff(HASH_LEN_IN_BYTES);
    this->computeHash(buff);
 
#if DEBUG_HASH
    std::cerr << "--- copmuted hash: 0x";
    for (unsigned i = 0; i < HashDecorator::HASH_LEN_IN_BYTES; ++i)
        std::cerr << std::hex << std::setw(2)
            << static_cast<unsigned int>(buff[i]);
    std::cerr << std::endl;
#endif
 
    // compare current hash with the stored one
    return hash == buff;
}
 
/**
 * Write data to image using ECC (Hamming) code, then store its length
 * and store its md5 hash to the image.
 * @param image QImage instance to store data to.
 * @param str Standard input stream to read data from.
 */
void writeData(QImage &image, std::istream &str) {
    ImageBlkIO gate(image);
    HamDecorator ham(gate);
    HashDecorator hash(ham);
    LenDecorator len(hash);
 
    // generic reference invariant to decorators' order
    IBlkIO &io = len;
 
    size_t inputLength = 0;
    const size_t blkSize = io.blkSize();
 
    // for each block
    for (unsigned i = 0; i < io.blkCnt(); ++i) {
        TData data(blkSize);
 
        // for each byte of the block
        for (unsigned j = 0; j < blkSize; ++j) {
            char c = 0;
            if (str.get(c))
                ++inputLength;
 
            data[j] = c;
        }
 
        // write block
        io.write(i, data);
    }
 
    if (str) {
        std::cerr << Color(C_LIGHT_RED)
            << "--- input exceeds maximal length which is "
            << blkSize * io.blkCnt() << Color(C_NO_COLOR)
            << std::endl;
    }
 
    len.setLength(inputLength);
    hash.wrtHash();
}
 
/**
 * Read data from image using ECC (Hamming) code and check its md5 hash.
 * @param image QImage instance to read data from.
 * @param str Standard output stream to write data to.
 */
void readData(QImage &image, std::ostream &str) {
    ImageBlkIO gate(image);
    HamDecorator ham(gate);
    HashDecorator hash(ham);
    LenDecorator len(hash);
 
    // generic reference invariant to decorators' order
    IBlkIO &io = len;
 
    const size_t inputLength = len.getLength();
    const size_t blkSize = io.blkSize();
 
#if DEBUG_LEN
    std::cerr << "--- blkSize:\t" << blkSize << std::endl;
    std::cerr << "--- blkCnt:\t" << io.blkCnt() << std::endl;
    std::cerr << "--- length:\t" << len.getLength() << std::endl;
#endif
    if (blkSize * io.blkCnt() < len.getLength()) {
        std::cerr << Color(C_LIGHT_RED)
            << "--- stored input length is not valid, will be ignored"
            << Color(C_NO_COLOR) << std::endl;
    }
 
    // for each block
    size_t pos = 0;
    for (unsigned i = 0; pos < inputLength && i < io.blkCnt(); ++i) {
        TData data(blkSize);
        io.read(i, data);
 
        // for each byte of the block
        for (unsigned j = 0; pos < inputLength && j < blkSize; ++j, ++pos)
            str << data[j];
    }
 
    if (ham.hasError()) {
        std::cerr << Color(C_LIGHT_RED)
            << "--- ECC error detected"
            << Color(C_NO_COLOR) << std::endl;
    }
 
    if (!hash.chkHash()) {
        std::cerr << Color(C_LIGHT_RED)
            << "--- hash error detected"
            << Color(C_NO_COLOR) << std::endl;
    }
}
 
namespace {
    void usage(const char *name) {
        std::cerr << "Usage:" << std::endl
            << "    " << name << " INPUT_IMAGE OUTPUT_IMAGE    "
            << "read INPUT_IMAGE, store data from stdin and save as OUTPUT_IMAGE"
            << std::endl
            << "    " << name << " INPUT_IMAGE                 "
            << "read INPUT_IMAGE and dump stored data to stdout"
            << std::endl;
    }
}
 
int main(int argc, char *argv[]) {
    const char *szAppName = argv[0];
    Color::enable(isatty(STDERR_FILENO));
 
    if (3 < argc || argc < 2) {
        usage(szAppName);
        return 1;
    }
 
    const char *fileIn = argv[1];
    const char *fileOut = 0;
 
    // tool's mode
    // true means "write data to image"
    // false means "read data from image"
    bool wrtMode;
    if ((wrtMode = (argc == 3)))
        fileOut = argv[2];
 
    try {
        // read image from input file (in both modes)
        QImage image(fileIn);
        if (image.isNull())
            throw std::runtime_error("can't read image from input file");
 
        if (wrtMode)
            writeData(image, std::cin);
        else
            readData(image, std::cout);
 
        // write image to output file (wrtMode only)
        // TODO: support other file formats
        if (wrtMode && !image.save(fileOut, "PNG"))
            throw std::runtime_error("can't write image to output file");
    }
    catch (std::runtime_error &e) {
        std::cerr << szAppName << ": "
            << Color(C_LIGHT_RED) << "error: " << Color(C_NO_COLOR)
            << e.what() << std::endl;
        return 1;
    }
 
    return 0;
}