/* pnmstega.c - insert a steganographic message into a portable anymap
**
** Copyright (C) 1994 by Jef Poskanzer.
**
** Permission to use, copy, modify, and distribute this software and its
** documentation for any purpose and without fee is hereby granted, provided
** that the above copyright notice appear in all copies and that both that
** copyright notice and this permission notice appear in supporting
** documentation.  This software is provided "as is" without express or
** implied warranty.
*/

#include "pnm.h"
#include <sys/types.h>
#include <sys/stat.h>

static void put_bit( int v, unsigned long loc );

static int comp_r, comp_g, comp_b, num_comps, comp_map[3];
static unsigned long isize;
static int format, cols;
static xel** xels;

#define LENLEN 32	/* length of the length field */

int
main( int argc, char* argv[] )
    {
    FILE* ifp;
    FILE* sfp;
    struct stat sb;
    long seed;
    unsigned long fsize, psize, plen, loc;
    char* cp;
    int argn, len, rows, byte, bitnum;
    xelval maxval;
    unsigned char c;
    static unsigned char bitmask[8] = {
	0x01, 0x02, 0x04, 0x08, 0x10, 0x20, 0x40, 0x80 };
    char* usage = "[-[no]length] [-seed n] [-components rgb] file [pnmfile]";

    pnm_init( &argc, argv );

    argn = 1;
    len = 0;
    seed = 0;
    comp_r = comp_g = comp_b = 1;
    num_comps = 3;

    while ( argn < argc && argv[argn][0] == '-' && argv[argn][1] != '\0' )
	{
	if ( pm_keymatch( argv[argn], "-length", 2 ) )
	    len = 1;
	else if ( pm_keymatch( argv[argn], "-nolength", 2 ) )
	    len = 0;
	else if ( pm_keymatch( argv[argn], "-seed", 2 ) )
	    {
	    ++argn;
	    if ( argn == argc || sscanf( argv[argn], "%ld", &seed ) != 1 )
		pm_usage( usage );
	    }
	else if ( pm_keymatch( argv[argn], "-components", 2 ) )
	    {
	    ++argn;
	    if ( argn == argc )
		pm_usage( usage );
	    comp_r = comp_g = comp_b = 0;
	    for ( cp = argv[argn]; *cp != '\0'; ++cp )
		switch ( *cp )
		    {
		    case 'r': case 'R': comp_r = 1; break;
		    case 'g': case 'G': comp_g = 1; break;
		    case 'b': case 'B': comp_b = 1; break;
		    default: pm_usage( usage );
		    }
	    num_comps = 0;
	    if ( comp_r ) ++num_comps;
	    if ( comp_g ) ++num_comps;
	    if ( comp_b ) ++num_comps;
	    }
	else
	    pm_usage( usage );
	++argn;
	}
    
    if ( seed < 0 )
	pm_error( "seed must be positive" );
    if ( seed != 0 )
	pm_srandom( seed );
    else if ( len )
	pm_message(
	    "warning, using the default seed with a length field is insecure" );

    if ( argn == argc )
	pm_usage( usage );
    sfp = pm_openr( argv[argn] );
    if ( fstat( fileno( sfp ), &sb ) < 0 )
	pm_perror( "fstat" );
    fsize = sb.st_size;
    plen = fsize * 8;
    if ( len )
	plen += LENLEN;
    ++argn;

    if ( argn != argc )
	{
	ifp = pm_openr( argv[argn] );
	++argn;
	}
    else
	ifp = stdin;

    if ( argn != argc )
	pm_usage( usage );

    xels = pnm_readpnm( ifp, &cols, &rows, &maxval, &format );
    pm_closer( ifp );
    if ( PNM_FORMAT_TYPE(format) == PBM_TYPE )
	pm_error( "can't stega into a PBM file - use PGM or PPM" );
    if ( PNM_FORMAT_TYPE(format) == PGM_TYPE )
	num_comps = 1;
    isize = rows * cols;
    psize = isize * num_comps;

    pm_message( "hiding %d bits in %d bit locations", plen, psize );
    if ( plen > psize )
	{
	if ( num_comps < 3 )
	    pm_error(
		"file won't fit, use more color components or a larger image" );
	else
	    pm_error( "file won't fit, use a larger image" );
	}

    if ( seed == 0 )
	loc = 0;
    else
	pm_init_perm( plen, psize );

    pm_map_comps( comp_r, comp_g, comp_b, num_comps, comp_map );

    if ( len )
	{
	for ( bitnum = 0; bitnum < LENLEN; ++bitnum )
	    {
	    if ( seed == 0 )
		put_bit( ( fsize >> bitnum ) & 0x1, loc++ );
	    else
		put_bit( ( fsize >> bitnum ) & 0x1, pm_next_perm() );
	    }
	}

    for ( byte = 0; byte < fsize; ++byte )
	{
	if ( fread( &c, 1, 1, sfp ) != 1 )
	    pm_error( "error reading file" );
	for ( bitnum = 0; bitnum < 8; ++bitnum )
	    {
	    if ( seed == 0 )
		put_bit( c & bitmask[bitnum], loc++ );
	    else
		put_bit( c & bitmask[bitnum], pm_next_perm() );
	    }
	}


    pnm_writepnm( stdout, xels, cols, rows, maxval, format, 0 );

    pm_closer( sfp );
    pm_closew( stdout );

    exit( 0 );
    }

static void
put_bit( int v, unsigned long loc )
    {
    register int comp, row, col;
    register xel* xP;
    register xelval r, g, b;

    comp = loc / isize;
    row = ( loc % isize ) / cols;
    col = loc % cols;
    xP = &(xels[row][col]);

    if ( PNM_FORMAT_TYPE(format) == PGM_TYPE )
	{
	g = PNM_GET1( *xP );
	if ( v )
	    g |= 1;
	else
	    g &= ~1;
	PNM_ASSIGN1( *xP, g );
	}
    else
	{
	r = PPM_GETR( *xP );
	g = PPM_GETG( *xP );
	b = PPM_GETB( *xP );
	if ( v )
	    {
	    switch ( comp_map[comp] )
		{
		case 0:
		r |= 1;
		break;
		case 1:
		g |= 1;
		break;
		case 2:
		b |= 1;
		break;
		default:
		pm_error( "bogus comp_map entry - shouldn't happen" );
		}
	    }
	else
	    {
	    switch ( comp_map[comp] )
		{
		case 0:
		r &= ~1;
		break;
		case 1:
		g &= ~1;
		break;
		case 2:
		b &= ~1;
		break;
		default:
		pm_error( "bogus comp_map entry - shouldn't happen" );
		}
	    }
	PPM_ASSIGN( *xP, r, g, b );
	}
    }