/* cryptimage.c */

#include <stdio.h>
#include <sys/types.h>

void usage() {
  printf("usage: cryptimage <imagefile> [messagefile]\n");
  printf("\nIf the \'messagefile\' param is provided, the message is hidden\n");
  printf("in 'imagefile' and output to a file called 'out.bmp'.  If \n");
  printf("\'messagefile\' is not provided, the program will try to extract\n");
  printf("a message from \'imagefile\' and write it to out.txt.\n\n");
  exit(1);
}


/*
 * Encode message into image. msb first
 *
 */
void encode(char *image,int imagelen,char *message,int messagelen) {
  int i,j,shift;
  int messagebits;
  long dataoff;
  FILE *c;
  unsigned char messagemask,imagemask;

  printf("Image of length %d\n",imagelen);
  printf("Message of length %d\n",messagelen);

  // calculate image size requirements to store
  // the message.  Try to use least # of bits possible
  for (messagebits=1;messagebits<5;messagebits++) {
    if (messagebits == 3)
      continue;
    if (((messagelen+4)*(8/messagebits)+2) < imagelen)
	break;
  }
  if (messagebits == 5) {
    printf("You need minimum %d bytes in your image to store your message\n",
      (messagelen+4)*(8/(messagebits)+2));
  }

  dataoff = *(long *)(image+10);
  printf("Image data starting at offset %d\n",dataoff);

  // first 2 bits represent number of bits per byte that are 
  // message bits.
  //
  //  00b - 1 bit
  //  01b - 2 bits
  //  11b - 4 bits
  //
  printf("Using %d bits per byte for message\n",messagebits);
  i=dataoff;
  image[i++] = (image[i]&0xFE) | (((messagebits-1))&0x01);
  image[i++] = (image[i]&0xFE) | (((messagebits-1)>>1)&0x01);

  imagemask = 0xFF << messagebits;
  messagemask = imagemask ^ 0xFF;

  // next 32 bits represents the message length
  printf("Storing message length of %d\n",messagelen);
  for (shift=(32-messagebits);shift>=0;i++) {
    image[i] = (image[i]&imagemask) | ((messagelen>>shift)&messagemask);
    shift -= messagebits;
  }
  
  printf("Hiding %d message bytes in %d image bytes\n",messagelen,messagelen*(8/messagebits));
  printf("Using message mask of 0x%02X\n",messagemask);
  printf("Using image mask of 0x%02X\n",imagemask);
  for (j=0,shift=(8-messagebits);j<messagelen;i++) {
    image[i] = (image[i]&imagemask) | ((message[j]>>shift)&messagemask);
    if (shift == 0) {
      shift = (8-messagebits);
      j++;
    }
    else
      shift-=messagebits;
  }

  printf("Opening file for output\n");
  c = fopen("out.bmp","w");
  fwrite(image,1,imagelen,c);
  fclose(c);
}

void decode(char *image,int imagelen) {
  int i,j,shift;
  long dataoff;
  char *message;
  unsigned long messagelen = 0;
  int  messagebits = 0;
  FILE *m;
  unsigned char messagemask,imagemask;

  printf("Image of length %d\n",imagelen);

  dataoff = *(long *)(image+10);
  printf("Image data starting at offset %d\n",dataoff);

  i = dataoff;
  messagebits |= image[i++]&0x01;
  messagebits <<=1;
  messagebits |= image[i++]&0x01;
  messagebits++;
  printf("Using %d bits per image byte for message\n",messagebits);
  imagemask = 0xFF << messagebits;
  messagemask = imagemask ^ 0xFF;
  printf("Using message mask of 0x%02X\n",messagemask);
  printf("Using image mask of 0x%02X\n",imagemask);

  for (shift=(32-messagebits);shift>=0;i++) {
    messagelen |= (image[i]&messagemask)<<shift;
    shift -= messagebits;
  }
  printf("Message length of %d\n",messagelen);
  
  message = (char *)malloc(messagelen);
  memset(message,0,messagelen);
  for (j=0,shift=(8-messagebits);j<messagelen;i++) {
    message[j] = message[j] | ((image[i]&messagemask)<<shift);
    if (shift == 0) {
      shift=(8-messagebits);
      j++;
    }
    else
      shift-=messagebits;
  }

  m = fopen("message.out","w");
  fwrite(message,1,messagelen,m);
  fclose(m);
}


int main(int argc,char **argv) {
  FILE *file;
  char *image=0,*message=0;
  int imagelen = 0,messagelen = 0;
  int bytesread;

  if ((argc != 2) && (argc != 3))
    usage();
  
  printf("Opening image file %s...",argv[1]);
  file = fopen(argv[1],"r");
  fseek(file,0,SEEK_END);
  imagelen = ftell(file);
  fseek(file,0,SEEK_SET);
  image = (char *)malloc(imagelen);
  bytesread = fread(image,1,imagelen,file);
  printf("read %d bytes.\n",bytesread);
  fclose(file);

  if (argc != 3)
    decode(image,imagelen);
  else {
    printf("Opening message file %s...",argv[2]);
    file = fopen(argv[2],"r");
    fseek(file,0,SEEK_END);
    messagelen = ftell(file);
    fseek(file,0,SEEK_SET);
    message = (char *)malloc(messagelen);
    bytesread = fread(message,1,messagelen,file);
    printf("read %d bytes.\n",bytesread);
    fclose(file);
    encode(image,imagelen,message,messagelen);
  }

  free(image);
  free(message);

  return 0;
}