// Copyright 2003 Nokia Corporation.
//
// THIS SOURCE CODE IS PROVIDED 'AS IS', WITH NO WARRANTIES WHATSOEVER,
// EXPRESS OR IMPLIED, INCLUDING ANY WARRANTY OF MERCHANTABILITY, FITNESS
// FOR ANY PARTICULAR PURPOSE, OR ARISING FROM A COURSE OF DEALING, USAGE
// OR TRADE PRACTICE, RELATING TO THE SOURCE CODE OR ANY WARRANTY OTHERWISE
// ARISING OUT OF ANY PROPOSAL, SPECIFICATION, OR SAMPLE AND WITH NO
// OBLIGATION OF NOKIA TO PROVIDE THE LICENSEE WITH ANY MAINTENANCE OR
// SUPPORT. FURTHERMORE, NOKIA MAKES NO WARRANTY THAT EXERCISE OF THE
// RIGHTS GRANTED HEREUNDER DOES NOT INFRINGE OR MAY NOT CAUSE INFRINGEMENT
// OF ANY PATENT OR OTHER INTELLECTUAL PROPERTY RIGHTS OWNED OR CONTROLLED
// BY THIRD PARTIES
//
// Furthermore, information provided in this source code is preliminary,
// and may be changed substantially prior to final release. Nokia Corporation
// retains the right to make changes to this source code at
// any time, without notice. This source code is provided for informational
// purposes only.
//
// Nokia and Nokia Connecting People are registered trademarks of Nokia
// Corporation.
// Java and all Java-based marks are trademarks or registered trademarks of
// Sun Microsystems, Inc.
// Other product and company names mentioned herein may be trademarks or
// trade names of their respective owners.
//
// A non-exclusive, non-transferable, worldwide, limited license is hereby
// granted to the Licensee to download, print, reproduce and modify the
// source code. The licensee has the right to market, sell, distribute and
// make available the source code in original or modified form only when
// incorporated into the programs developed by the Licensee. No other
// license, express or implied, by estoppel or otherwise, to any other
// intellectual property rights is granted herein.
import org.bouncycastle.crypto.*;
import org.bouncycastle.crypto.digests.*;
import org.bouncycastle.crypto.engines.*;
import org.bouncycastle.crypto.modes.*;
import org.bouncycastle.crypto.paddings.*;
import org.bouncycastle.crypto.params.*;
import java.io.*;

// The class MessageCodec encodes and decodes messages
class MessageCodec
{
  final static int OPTION_DIGEST = 0x1;


  // Creates a BlockCipher Engine, for example DES or AES
  static BlockCipher createEngine()
  {
    return new CBCBlockCipher(new DESEngine());
  }

  // Creates a Digest Engine, for example MD5 or SHA1
  static Digest createDigest()
  {
    return new MD5Digest();
  }


  // Do the encoding of the message including the headers,
  // encrypted content and digest
  static byte[] encodeMessage(String plainText,
                              String password,
                              boolean addDigest)
    throws Exception
  {
    byte content[] = plainText.getBytes();
    byte key[] = password.getBytes();

    // Create the cipher. Modify createEngine to use another Engine
    BufferedBlockCipher cipherEngine =
        new PaddedBufferedBlockCipher(createEngine());
    // Initialize the cipher for encryption
    cipherEngine.init(true, new KeyParameter(key));

    byte[] cipherText = new byte[cipherEngine.getOutputSize(content.length)];
    byte[] digest = null;

    // Do encryption
    int cipherTextLength = cipherEngine.processBytes(content,
                                                     0,
                                                     content.length,
                                                     cipherText,
                                                     0);
    cipherEngine.doFinal(cipherText, cipherTextLength);

    // add a digest if required
    if (addDigest)
    {
      Digest digestEngine = createDigest();
      int digestSize = digestEngine.getDigestSize();
      digest = new byte[digestSize];
      digestEngine.update(content, 0, content.length);
      digestEngine.doFinal(digest, 0);
    }

    // Create temporary streams
    ByteArrayOutputStream out = new ByteArrayOutputStream();
    DataOutputStream dout = new DataOutputStream(out);

    // write length
    dout.writeShort(cipherText.length);
    // write options
    int options = (addDigest) ? OPTION_DIGEST : 0;

    dout.writeShort(options);
    // write cipher text
    out.write(cipherText);
    // write digest
    if (addDigest)
    {
      out.write(digest);
    }
    return out.toByteArray();
  }


  // Decodes a message reading the header,
  // decrypting the content and veryfing the password
  static String decodeMessage(byte[] messageContent, String password)
    throws Exception
  {
    byte key[] = password.getBytes();

    // create utility streams
    ByteArrayInputStream in = new ByteArrayInputStream(messageContent);
    DataInputStream din = new DataInputStream(in);

    // read the message content
    int cipherTextLength = din.readShort();
    int options = din.readShort();

    byte[] cipherText = new byte[cipherTextLength];
    in.read(cipherText, 0, cipherTextLength);

    byte[] digest = null;
    Digest digestEngine = createDigest();
    // If the message contains a digest read it
    if ((options & OPTION_DIGEST) == 1)
    {
      int digestSize = digestEngine.getDigestSize();
      digest = new byte[digestSize];
      in.read(digest, 0, digestSize);
    }

    // Create decryption engine
    BufferedBlockCipher cipherEngine =
        new PaddedBufferedBlockCipher(createEngine());
    // Initialize the cipher for decryption
    cipherEngine.init(false, new KeyParameter(password.getBytes()));

    byte[] plainText = new byte[cipherEngine.getOutputSize(cipherTextLength)];

    // do encryption
    int size = cipherEngine.processBytes(cipherText,
                                         0,
                                         cipherTextLength,
                                         plainText,
                                         0);
    cipherEngine.doFinal(plainText, size);

    // It is important to do trimming to remove possible padding
    String resultText = new String(plainText).trim();

    // verify the text digest
    if ((options & OPTION_DIGEST) == 1)
    {
      // calculate the digest of the resulting text
      plainText = resultText.getBytes();
      digestEngine.update(plainText, 0, plainText.length);
      byte[] calculatedDigest = new byte[digestEngine.getDigestSize()];

      digestEngine.doFinal(calculatedDigest, 0);
      // compare digests size
      if (calculatedDigest.length != digest.length)
      {
        throw new Exception("Digest size mismatch");
      }
      // compare byte per byte
      for (int i = 0; i < calculatedDigest.length; i++)
      {
        if (calculatedDigest[i] != digest[i])
        {
          throw new Exception("Digest misimatch. Integrity compromised");
        }
      }
    }
    return resultText;
  }

}