#include <cstdio>
#include <cstring>
#ifdef _WIN32
#include <windows.h>
#endif
#include <stdint.h>

#include "../Lib/aes/AES.h"
#include "../Lib/aes/AESUtil.h"
#include "../Lib/zlib/zutil.h"

#define BLOCKLEN 16
#define MAXDATA (16*BLOCKLEN) // 256 bytes -- max length of data
#define CRCBYTES 4
#define BLOCKMODE AES::CBC

void show_usage()
{
    fprintf( stderr, "Usage: sbd [-d] [-o outFilename] key ([-i inFilename]|[data]) \n"
             "  Use option -d to decrypt.\n"
             "  Options may appear in any order.\n"
             "  If key is exactly 32, 48, or 64 bytes, and consists of\n"
             "  the characters 0-9, A-F, & a-f, it is interpretes as hex.\n"
             "  If the key is <=16 chars, it is expanded to a 128-bit key.\n"
             "  17 to 24 chars are expanded to a 192-bit key, while > 24\n"
             "  chars are expanded to a 256-bit key.\n"
             "  The key must appear before data if data is included.\n"
             "  Data may be piped in.\n"
             "  Unencoded data max size is 256 bytes, may be supplied via\n"
             "  stdin, file, or on command line\n" );
}

int main( int argc, char **argv )
{
    if( argc == 1 )
    {
        show_usage();
        return 1;
    }

    bool decrypt( false );
    char* inFilename( NULL );
    char* outFilename( NULL );
    char* keyText( NULL );
    char* data( NULL );

    // Let's read the input parameters
    for( int i = 1; i < argc; ++i )
    {
        char* arg = argv[i];
        if( strlen( arg ) == 2 && arg[0] == '-' && arg[1] == 'd' )
        {
            decrypt = true;
        }
        else if( strlen( arg ) == 2 && arg[0] == '-' && arg[1] == 'o' )
        {
            if( i + 1 < argc )
            {
                outFilename = argv[++i];
            }
        }
        else if( strlen( arg ) == 2 && arg[0] == '-' && arg[1] == 'i' )
        {
            if( i + 1 < argc )
            {
                inFilename = argv[++i];
            }
        }
        else if( NULL == keyText )
        {
            keyText = argv[i];
        }
        else
        {
            if( NULL != data || NULL != inFilename )
            {
                show_usage();
                return 1;
            }
            data = argv[i];
        }
    }

    unsigned char inBytes[MAXDATA + CRCBYTES];
    unsigned char outBytes[MAXDATA + CRCBYTES];
    unsigned char testBytes[MAXDATA + CRCBYTES];
    int dataLen( 0 );

    int maxRead = MAXDATA + ( decrypt ? CRCBYTES : 0 );

    // Fill up the inBytes buffer
    if( NULL == data )
    {
        FILE* inFile( stdin );
        if( NULL != inFilename )
        {
            inFile = fopen( inFilename, "rb" );
            if( NULL == inFile )
            {
                fprintf( stderr, "%s could not be read.\n", inFilename );
                return 1;
            }
        }
        while( dataLen < maxRead && !feof( inFile ) )
        {
            int c = getc( inFile );
            if( c < 0 )
            {
                break;
            }
            inBytes[dataLen++] = c;
        }
    }
    else
    {
        dataLen = strlen( data );
        if( dataLen > maxRead ) dataLen = maxRead;
        for( int i = 0; i < dataLen; ++i )
        {
            inBytes[i] = data[i];
        }
    }

    // get the stored crc;
    uint32_t crc( 0 );
    if( decrypt )
    {
        if( dataLen < 4 )
        {
            fprintf( stderr, "Not enough data to decrypt!" );
            return 1;
        }
        for( int i = 0; i < 4; ++i )
        {
            crc <<= 8;
            crc |= ( int )( unsigned char )inBytes[dataLen - 1];
            dataLen--;
        }
    }

    // pad the data...
    int blocks = ( dataLen + BLOCKLEN - 1 ) / BLOCKLEN;
    int padDataLen = BLOCKLEN * blocks;
    for( int i = dataLen; i < padDataLen; ++i )
    {
        inBytes[i] = 0;
    }

    // Deal with the key
    if( NULL == keyText )
    {
        show_usage();
        return 1;
    }

    unsigned char key[32];
    int keyLen = makeKey( keyText, key );
    /*
    printf( "key:" );
    for ( int i = 0; i < keyLen; ++i )
    {
        printf( "%02X", key[i] );
    }
    printf( "\n\n" );
    */

    // Do the encryption
    AES crypt;
    crypt.SetParameters( keyLen * 8, BLOCKLEN * 8 );
    if( decrypt )
    {
        crypt.StartDecryption( key );
        //fprintf(stderr,"Decrypt dataLen=%d, key=%s, bytes=%d, blocks=%d\n", dataLen, key, *inBytes, blocks);
        crypt.Decrypt( inBytes, outBytes, blocks, BLOCKMODE );
        uint32_t crcChk = crc32( 0, outBytes, dataLen );
        if( crcChk != crc )
        {
            fprintf( stderr, "CRC Error 0x%08lX != computed 0x%08lX\n", ( unsigned long )crc, ( unsigned long )crcChk );
            //return 1;
        }
    }
    else
    {
        crypt.StartEncryption( key );
        crypt.Encrypt( inBytes, outBytes, blocks, BLOCKMODE );
        crypt.StartDecryption( key );
        //fprintf(stderr,"Decrypt dataLen=%d, key=%s, bytes=%d, blocks=%d\n", dataLen, key, *outBytes, blocks);
        crypt.Decrypt( outBytes, testBytes, blocks, BLOCKMODE );
        if( memcmp( inBytes, testBytes, dataLen ) != 0 )
        {
            fprintf( stderr, "Error: decryption test error\n" );
            return 1;
        }
        // Add the crc
        crc = crc32( 0, inBytes, padDataLen );
        for( int i = 0; i < 4; ++i )
        {
            outBytes[padDataLen++] = crc & 0xFF;
            crc >>= 8;
        }
    }

    FILE* outFile( stdout );
    if( NULL != outFilename )
    {
        outFile = fopen( outFilename, "wb" );
        if( NULL == outFile )
        {
            fprintf( stderr, "%s could not be written.\n", outFilename );
            return 1;
        }
    }

    for( int i = 0; i < padDataLen; ++i )
    {
        if( fputc( outBytes[i], outFile ) == EOF )
        {
            fprintf( stderr, "Error writing to %s\n", outFilename );
            return 1;
        }
    }

    if( NULL != outFilename )
    {
        fclose( outFile );
    }

    return 0;
}