//---------------------------------------------------------------------------

// JTHZ steganographic example (Win32 console mode program).
// Borland C++Builder 5 PRO patch #1.

// Embeds a file in a 24 bit Windows DIB bitmap (*.bmp).
// example e bitmap.bmp inputfile

// Extracts a file from a 24 bit Windows DIB bitmap (*.bmp).
// example x bitmap.bmp outputfile

// The least significant bit of all 3 bytes (RGB) of every pixel
// can be used. First we embed the filesize, then the entire file:
// [FILESIZE][FILEDATA][n BYTES OF ORIGINAL UNUSED DATA]
// ...where n can be 0.

//---------------------------------------------------------------------------

// Allow iostream output for AnsiString type.
#define VCL_IOSTREAM

#include <vcl.h>
#pragma hdrstop

// #define NDEBUG to turn off assertions.
#define NDEBUG

#include <algorithm>
#include <assert.h>
#include <bitset>
#include <iostream>
#include <memory>
#include <vector>

//---------------------------------------------------------------------------

// Prototypes.

int Usage();
bool IsRgbDib(Graphics::TBitmap*);
int Embed(const AnsiString, const AnsiString);
int Extract(const AnsiString, const AnsiString);

//---------------------------------------------------------------------------

// Constants.

static const int CBadFileHandle   = -1;
static const int CBitsInOneByte   = 8;
static const int CBytesInOnePixel = 3;
static const int CFindSuccess     = 0;
static const int CReturnError     = -1;
static const int CReturnOK        = 0;

//---------------------------------------------------------------------------

int main(int argc, char* argv[])
{

  // We need a 4 argument commandline to run.
  if(argc != 4)
  {
    return Usage();
  }

  AnsiString Mode = LowerCase(argv[1]);
  assert(Mode.Length() > 0);

  // Pascalian AnsiString is 1-based, so:
  switch(Mode[1])
  {
    case 'e':
      return Embed(argv[3], argv[2]);

    case 'x':
      return Extract(argv[3], argv[2]);

    default:
      return Usage();
  }

}

//---------------------------------------------------------------------------

// Display "manual" of this program.

int Usage()
{

  std::cout << "usage:"
            << std::endl
            << "example e bmpfile inputfile"
            << std::endl
            << "example x bmpfile outputfile"
            << std::endl;

  return CReturnError;

}

//---------------------------------------------------------------------------

// Returns true if bmp is a 24 bit RGB DIB bitmap.
// As you can see, this code relies heavily on the VCL
// Graphics::TBitmap wrapper class.

bool IsRgbDib(Graphics::TBitmap* Bitmap)
{

  assert(Bitmap);

  return (Bitmap->PixelFormat == pf24bit) &&
         (Bitmap->HandleType  == bmDIB);

}

//---------------------------------------------------------------------------

// Try to embed inputFile in bitmap file.

int Embed
  (
    const AnsiString InputFileName,
    const AnsiString BitmapFileName
  )
{

  // Check if bitmap file exists at all.
  if(!FileExists(BitmapFileName))
  {
    std::cout << BitmapFileName
              << " does not exists"
              << std::endl;

    return CReturnError;
  }

  // Check if input file exists at all.
  if(!FileExists(InputFileName))
  {
    std::cout << InputFileName
              << " does not exists"
              << std::endl;

    return CReturnError;
  }

  // Figure out filesize of input file.

  TSearchRec SearchRec;

  int FindResult = FindFirst(InputFileName, faAnyFile, SearchRec);

  if(FindResult != CFindSuccess)
  {
    std::cout << "error getting filesize of "
              << InputFileName
              << std::endl;

    return CReturnError;
  }

  int FileSize = SearchRec.Size;

  FindClose(SearchRec);

  // Instantiate a TBitmap object. This object will load the
  // original BMP, 'contaminate' it with steganographic data
  // and save it back to disk.
  std::auto_ptr <Graphics::TBitmap> Bitmap(new Graphics::TBitmap);

  assert(Bitmap.get());

  try
  {
    Bitmap->LoadFromFile(BitmapFileName);
  }
  catch(...)
  {
    std::cout << "error reading "
              << BitmapFileName
              << " as DIB"
              << std::endl;

    return CReturnError;
  }

  if(!IsRgbDib(Bitmap.get()))
  {
    std::cout << BitmapFileName
              << ": not a 24 bit RGB DIB bitmap"
              << std::endl;

    return CReturnError;
  }

  // You can only squeeze so much data into a given bitmap.
  // If so desired, you can use more bits (instead of just 1).
  // In our JTHZ Contraband Hell Edition, we use a variable
  // amount of bits, which is a very cool concept. Turns out you
  // can easily use 50% or more of all bits w/o too much
  // damage to the picture.
  int UsefulBits = Bitmap->Height * Bitmap->Width * CBytesInOnePixel;

  int NeededBits = CBitsInOneByte * (sizeof(int) + FileSize);

  if(NeededBits > UsefulBits)
  {
    std::cout << "inputfile too large to embed in bitmap"
              << std::endl;

    return CReturnError;
  }

  // Create a buffer that contains all the data we will embed
  // i.e. the size of the file and the filedata.
  int EmbedDataSize = sizeof(int) + FileSize;
  std::vector <Byte> EmbedData(EmbedDataSize);

  // Copy fileSize to first 4 bytes.
  memcpy(EmbedData.begin(), &FileSize, sizeof(int));

  // Read entire file into the next part of EmbedData.
  int InputFile = FileOpen(InputFileName, fmOpenRead);

  if(InputFile == CBadFileHandle)
  {
    std::cout << "error opening inputfile"
              << std::endl;

    return CReturnError;
  }

  int BytesRead = FileRead
  (
    InputFile,
    EmbedData.begin() + sizeof(int),
    FileSize
  );

  FileClose(InputFile);

  if(BytesRead != FileSize)
  {
    // Oddball error; pro'lly a defect bitmap or disk/file error.
    std::cout << "error reading "
              << InputFileName
              << std::endl;

    return CReturnError;
  }

  // Embed data into bitmap.
  int Line           = 0;
  int ByteOffset     = 0;

  int BytesInOneLine = CBytesInOnePixel * Bitmap->Width;

  // In this loop, we iterate the data we want to embed and
  // update the scanline and byte-offset into the bitmap on the fly.
  // You can also turn this around and iterate your bitmap until
  // you're done embedding, but that would yield two for-loops with
  // multiple exit-points, wouldn't you think?
  for(int Idx = 0; Idx < EmbedDataSize; ++Idx)
  {
    for(Byte BitIdx = 0; BitIdx < CBitsInOneByte; ++BitIdx)
    {

      Byte* ScanlinePtr = reinterpret_cast <Byte*>
                          (Bitmap->ScanLine[Line]);

      // Clear least significant bit.
      ScanlinePtr[ByteOffset] &= 0xfe;

      // 'OR' in the next bit from our data.
      Byte NextBit = static_cast <Byte>
                     ((EmbedData.at(Idx) >> BitIdx) & 0x01);

      ScanlinePtr[ByteOffset] |= NextBit;

      // Move to next byte.
      ++ByteOffset;

      if(ByteOffset >= BytesInOneLine)
      {
        // Move to next scanline.

        ByteOffset = 0;

        ++Line;
      }
    }
  }

  // Save data.
  try
  {
    Bitmap->SaveToFile(BitmapFileName);
  }
  catch(...)
  {
    std::cout << "error saving to "
              << BitmapFileName
              << std::endl;

    return CReturnError;
  }

  return CReturnOK;

}

//---------------------------------------------------------------------------

// Try to extract embedded data from bitmap file to output file.

int Extract
  (
    const AnsiString OutputFileName,
    const AnsiString BitmapFileName
  )
{

  // See if bitmap file exists at all.
  if(!FileExists(BitmapFileName))
  {
    std::cout << BitmapFileName
              << " does not exists"
              << std::endl;

    return CReturnError;
  }

  std::auto_ptr <Graphics::TBitmap> Bitmap(new Graphics::TBitmap);

  assert(Bitmap.get());

  try
  {
    Bitmap->LoadFromFile(BitmapFileName);
  }
  catch(...)
  {
    std::cout << "error reading "
              << BitmapFileName
              << " as DIB"
              << std::endl;

    return CReturnError;
  }

  if(!IsRgbDib(Bitmap.get()))
  {
    std::cout << BitmapFileName
              << ": not a 24 bit RGB DIB bitmap"
              << std::endl;

    return CReturnError;
  }

  // See if there is something embedded in the bitmap.

  int UsefulBits = Bitmap->Height * Bitmap->Width * CBytesInOnePixel;

  if(UsefulBits < ((sizeof(int) * CBitsInOneByte)))
  {
    std::cout << "No data embedded in "
              << BitmapFileName
              << std::endl;

    return CReturnError;
  }

  // Extract embedded fileSize from bitmap.
  int Line           = 0;
  int ByteOffset     = 0;

  int BytesInOneLine = CBytesInOnePixel * Bitmap->Width;

  // Retrieve FileSize as a series of bits.
  // This results in a more readable and shorter loop.
  std::bitset <sizeof(int) * CBitsInOneByte> FileSizeAsBitset;

  const int EmbeddedFileSizeBits = sizeof(int) * CBitsInOneByte;

  for(int Idx = 0; Idx < EmbeddedFileSizeBits; ++Idx)
  {
    Byte* ScanlinePtr = reinterpret_cast <Byte*>
                        (Bitmap->ScanLine[Line]);

    // If ScanlinePtr[ByteOffset] AND 1, this bit is part
    // of the embedded filesize.
    FileSizeAsBitset.set(Idx, ScanlinePtr[ByteOffset] & 0x01);

    // Move to next byte.
    ++ByteOffset;

    if(ByteOffset >= BytesInOneLine)
    {
      // Move to next scanline.

      ByteOffset = 0;

      ++Line;
    }
  }

  // Perform cast because std::bitset has no to_int() method.
  int FileSize = static_cast <int>
                 (FileSizeAsBitset.to_ulong());

  // We're gonna keep the values of ByteOffset and Line,
  // because we have to continue reading from there.

  int MaxDataBits = UsefulBits - (sizeof(int) * CBitsInOneByte);

  if((FileSize * CBitsInOneByte) > MaxDataBits)
  {
    std::cout << "no data embedded in "
              << BitmapFileName
              << std::endl;

    return CReturnError;
  }

  // Extract data from bitmap.

  std::vector <Byte> ExtractData(FileSize);

  std::fill(ExtractData.begin(), ExtractData.end(), 0);

  // Again, we're looping until we have all the data that's
  // embedded, updating indices into the bitmap (scanline
  // and byte offset) as we go along.
  for(int ByteIdx = 0; ByteIdx < FileSize; ++ByteIdx)
  {

    for(Byte BitIdx = 0; BitIdx < CBitsInOneByte; ++BitIdx)
    {

      Byte* ScanlinePtr = reinterpret_cast <Byte*>
                          (Bitmap->ScanLine[Line]);

      // Retrieve next bit and shift it in.
      ExtractData[ByteIdx] |= static_cast <Byte>
                              ((ScanlinePtr[ByteOffset] & 0x01) << BitIdx);

      // Move to next byte.
      ++ByteOffset;

      if(ByteOffset >= BytesInOneLine)
      {
        // Move to next scanline.

        ByteOffset = 0;

        ++Line;
      }
    }
  }

  // Save to file.

  int OutputFile = FileCreate(OutputFileName);

  if(OutputFile == CBadFileHandle)
  {
    std::cout << "error (re)creating "
              << OutputFileName
              << std::endl;

    return CReturnError;
  }

  int BytesWritten = FileWrite
  (
    OutputFile,
    ExtractData.begin(),
    FileSize
  );

  FileClose(OutputFile);

  if(BytesWritten != FileSize)
  {
    std::cout << "error writing to "
              << OutputFileName
              << std::endl;

    return CReturnError;
  }

  return CReturnOK;

}

//---------------------------------------------------------------------------