/*
 * Copyright (c) 2010, Manuel Mausz. All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions
 * are met:
 *
 *   - Redistributions of source code must retain the above copyright
 *     notice, this list of conditions and the following disclaimer.
 *
 *   - Redistributions in binary form must reproduce the above copyright
 *     notice, this list of conditions and the following disclaimer in the
 *     documentation and/or other materials provided with the distribution.
 *
 *   - The names of the authors may not be used to endorse or promote products
 *     derived from this software without specific prior written permission.
 *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS
 * IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO,
 * THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
 * PURPOSE ARE DISCLAIMED.  IN NO EVENT SHALL THE COPYRIGHT OWNER OR
 * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
 * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
 * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
 * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF
 * LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING
 * NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
 * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 */

import java.io.*;
import javax.crypto.*;
import java.util.Collection;
import java.util.Iterator;
import org.bouncycastle.util.encoders.Base64;
import java.security.GeneralSecurityException;

/*
 * Common static methods used by all applications
 *
 * @author Manuel Mausz (0728348)
 */
class Utils
{
  static class Shutdown
    extends java.lang.RuntimeException
  {
    public Shutdown()
    {
      super();
    }

    public Shutdown(String message)
    {
      super(message);
    }
  }

  /*==========================================================================*/

  private static final String B64 = "a-zA-Z0-9/+";

  /*==========================================================================*/

  public static <T> String join(final Collection<T> objs, final String delimiter)
  {
    if (objs == null || objs.isEmpty())
      return "";
    Iterator<T> iter = objs.iterator();
    StringBuffer buffer = new StringBuffer(iter.next().toString());
    while (iter.hasNext())
      buffer.append(delimiter).append(iter.next().toString());
    return buffer.toString();
  }

  /*--------------------------------------------------------------------------*/

  public static long parseHeaderNum(String[] args, int index)
  {
    long num = -1;
    try
    {
      if (args.length < index + 1)
        throw new NumberFormatException();
      num = Long.parseLong(args[index]);
      if (num < 0)
        throw new NumberFormatException();
    }
    catch(NumberFormatException e)
    {
      System.err.println("Error: Invalid network paket from peer:");
      System.err.println(join(java.util.Arrays.asList(args), " "));
    }
    return num;
  }

  /*------------------------------------------------------------------------*/

  public static void sendOutput(EncObjectOutputStream out, String[] msgs)
    throws IOException
  {
    out.writeLine("!output " + msgs.length);
    for(String msg : msgs)
      out.writeLine("!" + msg);
    out.flush();
  }

  /*------------------------------------------------------------------------*/

  public static void sendOutput(EncObjectOutputStream out, String msg)
    throws IOException
  {
    String[] tmp = { msg };
    sendOutput(out, tmp);
  }

  /*------------------------------------------------------------------------*/

  public static void sendError(EncObjectOutputStream out, String[] msgs)
    throws IOException
  {
    out.writeLine("!error " + msgs.length);
    for(String msg : msgs)
      out.writeLine("!" + msg);
    out.flush();
  }

  /*------------------------------------------------------------------------*/

  public static void sendError(EncObjectOutputStream out, String msg)
    throws IOException
  {
    String[] tmp = { "Error: " + msg };
    sendError(out, tmp);
  }

  /*==========================================================================*/

  static class EncObjectInputStream
    extends ObjectInputStream
  {
    private Cipher cipher = null;
    private Mac mac = null;

    public EncObjectInputStream(InputStream in)
      throws IOException
    {
      super(in);
    }

    /*------------------------------------------------------------------------*/

    public void setCipher(Cipher cipher)
    {
      this.cipher = cipher;
    }

    /*------------------------------------------------------------------------*/

    public void setMAC(Mac mac)
    {
      this.mac = mac;
    }

    /*------------------------------------------------------------------------*/

    private String decode(String msg)
    {
      if (msg == null)
        return null;
      if ((cipher == null && mac == null) || msg.charAt(0) == '!')
        return msg;

      if (mac != null)
      {
        msg = new String(Base64.decode(msg));
        assert msg.matches("[" + B64 + "]{43}= [\\s[^\\s]]+");
        int ix = msg.indexOf(' ');
        if (ix == -1)
        {
          System.err.println("Error: invalid hashed message format:");
          System.err.println(msg);
          return null;
        }
        String hasht = msg.substring(0, ix);
        String msgtmp = msg.substring(ix + 1);
        String hashc = new String(Base64.encode(mac.doFinal(msgtmp.getBytes())));
        if (!hasht.equals(hashc))
        {
          System.err.println("Error: invalid MAC:");
          System.err.println(msg);
          throw new HashError();
          //return null;
        }
        msg = msgtmp;
      }

      if (cipher != null)
      {
        try
        {
          byte[] decmsg = Base64.decode(msg);
          return new String(cipher.doFinal(decmsg));
        }
        catch(GeneralSecurityException e)
        {
          System.err.println("Error: Unable to decode message: " + e.getMessage());
          return null;
        }
      }

      return msg;
    }

    /*------------------------------------------------------------------------*/

    public String readUTF()
      throws IOException
    {
      return decode(super.readUTF());
    }

    /*------------------------------------------------------------------------*/

    @SuppressWarnings("deprecation")
    public String readLine()
      throws IOException
    {
      return decode(super.readLine());
    }
  }

  /*==========================================================================*/

  static class EncObjectOutputStream
    extends ObjectOutputStream
  {
    private Cipher cipher = null;
    private Mac mac = null;

    public EncObjectOutputStream(OutputStream out)
      throws IOException
    {
      super(out);
    }

    /*------------------------------------------------------------------------*/

    public void setCipher(Cipher cipher)
    {
      this.cipher = cipher;
    }

    /*------------------------------------------------------------------------*/

    public void setMAC(Mac mac)
    {
      this.mac = mac;
    }

    /*------------------------------------------------------------------------*/

    private byte[] encode(byte[] msg)
    {
      try
      {
        return Base64.encode(cipher.doFinal(msg));
      }
      catch(GeneralSecurityException e)
      {
        System.err.println("Error: Unable to encode message: " + e.getMessage());
      }
      return null;
    }

    /*------------------------------------------------------------------------*/

    public void writeUTF(String msg)
      throws IOException
    {
      if (cipher == null && mac == null)
      {
        super.writeUTF(msg);
        return;
      }

      if (mac != null)
      {
        byte[] hash = mac.doFinal(msg.getBytes());
        msg = new String(Base64.encode(hash)) + " " + msg;
        assert msg.matches("[" + B64 + "]{43}= [\\s[^\\s]]+");
      }

      if (cipher != null)
      {
        super.writeUTF(new String(encode(msg.getBytes())));
        return;
      }

      String hashedmsg = new String(Base64.encode(msg.getBytes()));
      super.writeUTF(hashedmsg);
      return;
    }

    /*------------------------------------------------------------------------*/

    public void writeLine(String msg)
      throws IOException
    {
      if (cipher == null && mac == null)
      {
        super.writeBytes(msg + "\n");
        return;
      }

      if (mac != null)
      {
        byte[] hash = mac.doFinal(msg.getBytes());
        msg = new String(Base64.encode(hash)) + " " + msg;
        assert msg.matches("[" + B64 + "]{43}= [\\s[^\\s]]+");
      }

      if (cipher != null)
      {
        super.writeBytes(new String(encode(msg.getBytes())) + "\n");
        return;
      }

      super.writeBytes(new String(Base64.encode(msg.getBytes())) + "\n");
      return;
    }
  }

  /*==========================================================================*/

  static class HashError
    extends java.lang.RuntimeException
  {
    public HashError()
    {
      super();
    }
  }
}