/*
 * 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 static java.lang.System.err;
import static java.lang.System.out;

import java.util.Arrays;
import java.util.MissingResourceException;

import javax.crypto.*;
import java.security.*;
import org.bouncycastle.openssl.PEMReader;
import org.bouncycastle.util.encoders.Base64;
import org.bouncycastle.openssl.PasswordFinder;

import java.net.*;
import java.io.*;

/*
 * Client implementation for Lab#1 of DSLab WS10
 * See angabe.pdf for details
 *
 * This code is not documented at all. This is volitional
 *
 * @author Manuel Mausz (0728348)
 */
public class Client
{
  public class ProxyConnection
    extends CommandNetwork
    implements Runnable
  {
    private final Utils.EncObjectInputStream oin;
    private final Utils.EncObjectOutputStream oout;
    private final String sharedFolder;
    private final Object mainLock;
    private final Object intLock;

    ProxyConnection(Utils.EncObjectInputStream oin, Utils.EncObjectOutputStream oout,
        String sharedFolder, Object mainLock, Object intLock)
      throws NoSuchMethodException
    {
      this.oin          = oin;
      this.oout         = oout;
      this.sharedFolder = sharedFolder;
      this.mainLock     = mainLock;
      this.intLock      = intLock;

      cmdHandler.register("!error",    this, "cmdError");
      cmdHandler.register("!output",   this, "cmdOutput");
      cmdHandler.register("!download", this, "cmdDownload");
      cmdHandler.register("!upload",   this, "cmdUpload");
      cmdHandler.register("!ok",       this, "cmdOk");
      cmdHandler.register("unknown",   this, "cmdUnknown");
    }

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

    public void notifyInteractive()
    {
      synchronized(intLock)
      {
        intLock.notify();
      }
    }

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

    public void cmdPrintOutput(PrintStream out, String cmd, String[] args)
      throws IOException
    {
      long num;
      if ((num = Utils.parseHeaderNum(args, 0)) < 0)
        return;

      String msg;
      for (; num > 0 && (msg = oin.readLine()) != null; --num)
        out.println(msg.substring(1));
      notifyInteractive();
    }

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

    public void cmdError(String cmd, String[] args)
      throws IOException
    {
      cmdPrintOutput(err, cmd, args);
    }

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

    public void cmdOutput(String cmd, String[] args)
      throws IOException
    {
      cmdPrintOutput(out, cmd, args);
    }

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

    public void cmdDownload(String cmd, String[] args)
    {
      if (args.length < 2 || args[1].length() <= 0)
      {
        err.println("Error: Invalid " + cmd + "-command paket from proxy. Ignoring...");
        notifyInteractive();
        return;
      }

      long filesize;
      if ((filesize = Utils.parseHeaderNum(args, 1)) < 0)
        return;

      File file = new File(sharedFolder, args[0]);
      file.delete();
      try
      {
        FileOutputStream fout = null;
        try
        {
          fout = new FileOutputStream(file);
        }
        catch(FileNotFoundException e)
        {
          err.println("Error: Unable to write to file '" + file + "': " + e.getMessage());
        }

        byte[] buffer = new byte[8 * 1024];
        int toread = buffer.length;
        while(filesize > 0)
        {
          if (filesize < toread)
            toread = (int) filesize;
          int count = oin.read(buffer, 0, toread);
          if (count == -1)
            throw new IOException("Connection reset by peer");

          if (fout != null)
          {
            /* decode that chunk */
            byte[] decbuffer = new byte[aesdecrypt.getOutputSize(count)];
            int deccount = aesdecrypt.update(buffer, 0, count, decbuffer);
            fout.write(decbuffer, 0, deccount);
          }
          filesize -= count;
        }
        if (fout != null)
        {
          /* decryption must be finalized */
          byte[] decbuffer = aesdecrypt.doFinal();
          fout.write(decbuffer);
        }
      }
      catch(IOException e)
      {
        err.println("Error: Error during file transfer: " + e.getMessage());
      }
      catch(GeneralSecurityException e)
      {
        err.println("Error: Error during decrypting file transfer: " + e.getMessage());
      }

      notifyInteractive();
    }

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

    public void cmdUpload(String cmd, String[] args)
    {
      if (args.length < 1 || args[0].length() <= 0)
      {
        err.println("Error: Invalid " + cmd + "-command paket from proxy. Ignoring...");
        notifyInteractive();
        return;
      }

      File file = new File(sharedFolder, args[0]);
      if (!file.exists() || !file.getParent().equals(sharedFolder))
      {
        err.println("Error: File '" + args[0] + "' does not exist");
        notifyInteractive();
        return;
      }
      if (!file.isFile())
      {
        err.println("Error: File '" + args[0] + "' is not a file");
        notifyInteractive();
        return;
      }
      if (!file.canRead())
      {
        err.println("Error: File '" + args[0] + "' is not readable");
        notifyInteractive();
        return;
      }

      long filesize = file.length();
      try
      {
        FileInputStream fin = new FileInputStream(file);

        oout.writeLine("!upload2 " + file.getName() + " " + filesize);

        byte[] buffer = new byte[8 * 1024];
        int toread = buffer.length;
        while(filesize > 0)
        {
          if (filesize < toread)
            toread = (int) filesize;
          int count = fin.read(buffer, 0, toread);
          if (count == -1)
            throw new IOException("Error while reading from file");

          /* encode that chunk */
          byte[] encbuffer = new byte[aesencrypt.getOutputSize(count)];
          int enccount = aesencrypt.update(buffer, 0, count, encbuffer);

          oout.write(encbuffer, 0, enccount);
          filesize -= count;
        }
        /* encryption must be finalized */
        oout.write(aesencrypt.doFinal());
        oout.flush();
      }
      catch(FileNotFoundException e)
      {
        err.println("Error: File '" + args[0] + "' is not readable");
        notifyInteractive();
        return;
      }
      catch(IOException e)
      {
        err.println("Error during file transfer: " + e.getMessage());
        notifyInteractive();
        return;
      }
      catch(GeneralSecurityException e)
      {
        err.println("Error during encrypting file transfer: " + e.getMessage() + "");
        notifyInteractive();
        return;
      }
    }

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

    public void cmdOk(String cmd, String[] args)
      throws IOException
    {
      if (args.length != 4)
      {
        err.println("Error: Invalid " + cmd + "-command paket from proxy. Ignoring...");
        notifyInteractive();
        return;
      }

      if (!args[0].equals(mychallenge))
      {
        err.println("Error: Received invalid challenge from proxy. Ignoring...");
        notifyInteractive();
        return;
      }

      String secondMessage = cmd + " " + Utils.join(Arrays.asList(args), " ");
      assert secondMessage.matches("!ok [" + B64 + "]{43}= [" + B64 + "]{43}= [" + B64 + "]{43}= [" + B64 + "]{22}==") : "2nd message";

      String proxychallenge = args[1];
      javax.crypto.spec.SecretKeySpec seckey = new javax.crypto.spec.SecretKeySpec(Base64.decode(args[2].getBytes()), "AES");
      javax.crypto.spec.IvParameterSpec iv = new javax.crypto.spec.IvParameterSpec(Base64.decode(args[3].getBytes()));

      try
      {
        aesencrypt.init(Cipher.ENCRYPT_MODE, seckey, iv);
        aesdecrypt.init(Cipher.DECRYPT_MODE, seckey, iv);
      }
      catch(InvalidKeyException e)
      {
        err.println("Error: invalid AES key: " + e.getMessage());
        notifyInteractive();
        return;
      }
      catch(InvalidAlgorithmParameterException e)
      {
        err.println("Error: invalid AES parameter: " + e.getMessage());
        notifyInteractive();
        return;
      }

      oout.setCipher(aesencrypt);
      oin.setCipher(aesdecrypt);

      String thirdMessage = proxychallenge;
      assert thirdMessage.matches("[" + B64 + "]{43}=") : "3rd message"; 
      oout.writeLine(thirdMessage);
      oout.flush();
      loggedin = true;
    }

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

    public void cmdUnknown(String cmd, String[] args)
    {
      err.println("Error: Unknown data from proxy: " + cmd + " "
          + Utils.join(Arrays.asList(args), " "));
      notifyInteractive();
    }

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

    public void run()
    {
      try
      {
        run(oin);
      }
      catch(CommandHandler.Exception e)
      {
        err.println("Internal Error: " + e.getMessage());
      }
      catch(IOException e)
      {
        /* ignore that exception
         * thread will shutdown and unlock the main thread
         * which will shutdown the application
         */
      }

      notifyInteractive();
      synchronized(mainLock)
      {
        mainLock.notify();
      }
    }
  }

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

  public class Interactive
    extends CommandInteractive
    implements Runnable
  {
    private final InputStream sin;
    private final Utils.EncObjectInputStream oin;
    private final Utils.EncObjectOutputStream oout;
    private final Object mainLock;
    private final Object intLock;

    Interactive(InputStream sin, Utils.EncObjectInputStream oin,
        Utils.EncObjectOutputStream oout, Object mainLock, Object intLock)
      throws NoSuchMethodException
    {
      this.sin      = sin;
      this.mainLock = mainLock;
      this.intLock  = intLock;
      this.oin      = oin;
      this.oout     = oout;

      cmdHandler.register("unknown", this, "cmdUnknown");
      cmdHandler.register("!login",  this, "cmdLogin");
      cmdHandler.register("!exit",   this, "cmdExit");
    }

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

    public void waitForSocket()
    {
      synchronized(intLock)
      {
        try
        {
          intLock.wait(1000);
        }
        catch(InterruptedException e)
        {
          /* if we get interrupted -> ignore */
        }
      }
    }

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

    public void cmdUnknown(String cmd, String[] args)
      throws IOException
    {
      if (!loggedin)
      {
        err.println("Not logged in");
        return;
      }

      oout.writeLine(cmd + " " + Utils.join(Arrays.asList(args), " "));
      oout.flush();
      waitForSocket();
    }

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

    public void cmdLogin(String cmd, String[] args)
      throws IOException
    {
      if (args.length != 1)
      {
        err.println("Invalid Syntax: !login <username>");
        return;
      }
      if (loggedin)
      {
        err.println("Already logged in");
        return;
      }

      /* read users private key */
      final String user = args[0];
      File pemfile = new File(keysDir, user + ".pem");
      if (!pemfile.isFile())
      {
        err.println("Error: Private keyfile '" + pemfile + "' doesn't exist.");
        return;
      }
      if (!pemfile.canRead())
      {
        err.println("Error: Private keyfile '" + pemfile + "' is not readable.");
        return;
      }

      try
      {
        PEMReader in = new PEMReader(new FileReader(pemfile), new PasswordFinder()
            {
              @Override
              public char[] getPassword()
              {
                try
                {
                  /* reads the password from standard input for decrypting the private key */
                  out.println("Enter pass phrase: ");
                  return new BufferedReader(new InputStreamReader(sin)).readLine().toCharArray();
                }
                catch(IOException e)
                {
                  char[] tmp = {};
                  return tmp;
                }
              }
            });
        KeyPair keyPair = (KeyPair) in.readObject();
        privateKey = keyPair.getPrivate();
        rsadecrypt.init(Cipher.DECRYPT_MODE, privateKey);
        oin.setCipher(rsadecrypt);
      }
      catch(FileNotFoundException e)
      {
        err.println("Error while reading private key of " + user + ". Unable to read keyfile.");
        return;
      }
      catch(IOException e)
      {
        err.println("Error while reading private key of " + user + ". Maybe wrong pass phrase.");
        return;
      }
      catch(InvalidKeyException e)
      {
        err.println("Error: invalid key file: " + e.getMessage());
        return;
      }

      /* generates a 32 byte secure random number */
      SecureRandom secureRandom = new SecureRandom();
      byte[] tmp = new byte[32];
      secureRandom.nextBytes(tmp);
      mychallenge = new String(Base64.encode(tmp));

      String firstMessage = cmd + " " + Utils.join(Arrays.asList(args), " ") + " " + mychallenge;
      assert firstMessage.matches("!login \\w+ [" + B64 + "]{43}=") : "1st message";

      oout.writeLine(firstMessage);
      oout.flush();
      waitForSocket();
    }

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

    public void cmdExit(String cmd, String[] args)
    {
      stop();
    }

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

    public void printPrompt()
    {
      out.print(">: ");
      out.flush();
    }

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

    public void shutdown()
    {
      try
      {
        oout.flush();
      }
      catch(IOException e)
      {}
    }

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

    public void run()
    {
      try
      {
        run(sin);
      }
      catch(CommandHandler.Exception e)
      {
        err.println("Internal Error: " + e.getMessage());
      }
      catch(IOException e)
      {
        /* ignore that exception
         * thread will shutdown and unlock the main thread
         * which will shutdown the application
         */
      }

      shutdown();
      synchronized(mainLock)
      {
        mainLock.notify();
      }
    }
  }

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

  private static String sharedFolder;
  private static String proxyHost;
  private static int proxyTCPPort;
  private Socket sock = null;
  private InputStream sockin = null;
  private OutputStream sockout = null;
  private Utils.EncObjectInputStream oin = null;
  private Utils.EncObjectOutputStream oout = null;
  private Thread tPConnection = null;
  private Thread tInteractive = null;
  private InputStream stdin = null;
  private final Object interactiveLock = new Object();
  private final Object mainLock = new Object();

  private static String keysDir;
  private static String proxyKey;
  private String mychallenge;
  private PrivateKey privateKey = null;
  private Cipher rsaencrypt, rsadecrypt;
  private Cipher aesencrypt, aesdecrypt;
  private volatile boolean loggedin = false;
  private final String B64 = "a-zA-Z0-9/+";

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

  public Client()
  {
    try
    {
      rsaencrypt = Cipher.getInstance("RSA/NONE/OAEPWithSHA256AndMGF1Padding");
      rsadecrypt = Cipher.getInstance("RSA/NONE/OAEPWithSHA256AndMGF1Padding");
      aesencrypt = Cipher.getInstance("AES/CTR/NoPadding");
      aesdecrypt = Cipher.getInstance("AES/CTR/NoPadding");
    }
    catch(NoSuchAlgorithmException e)
    {
      bailout("Unable to initialize cipher: " + e.getMessage());
    }
    catch(NoSuchPaddingException e)
    {
      bailout("Unable to initialize cipher: " + e.getMessage());
    }
  }

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

  public static void usage()
    throws Utils.Shutdown
  {
    out.println("Usage: Client sharedFolder\n");
    out.println("sharedFolder\t...the directory to put downloaded files or to upload files from");

    // Java is some piece of crap which doesn't allow me to set exitcode w/o
    // using System.exit. Maybe someday Java will be a fully functional
    // programming language, but I wouldn't bet my money
    //System.exit(1);
    throw new Utils.Shutdown("FUCK YOU JAVA");
  }

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

  public void bailout(String error)
    throws Utils.Shutdown
  {
    err.println("Error: " + error);
    shutdown();

    // Java is some piece of crap which doesn't allow me to set exitcode w/o
    // using System.exit. Maybe someday Java will be a fully functional
    // programming language, but I wouldn't bet my money
    //System.exit(2);
    throw new Utils.Shutdown("FUCK YOU JAVA");
  }

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

  public void parseArgs(String[] args)
  {
    if (args.length != 1)
      usage();

    sharedFolder = args[0];
    File dldir = new File(sharedFolder);
    if (!dldir.isDirectory())
      bailout("sharedFolder '" + sharedFolder + "' is not a directory");
    if (!dldir.canWrite())
      bailout("sharedFolder '" + sharedFolder + "' is not writeable");
  }

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

  public void parseConfig()
  {
    Config config = null;
    try
    {
      config = new Config("client");
    }
    catch(MissingResourceException e)
    {
      bailout("configuration file doesn't exist or isn't readable");
    }

    String directive = null;
    try
    {
      directive = "proxy.host";
      proxyHost = config.getString(directive);
      if (proxyHost.length() == 0)
        bailout("configuration directive '" + directive + "' is empty or invalid");

      directive = "proxy.tcp.port";
      proxyTCPPort = config.getInt(directive);
      if (proxyTCPPort <= 0 || proxyTCPPort > 65536)
        bailout("configuration directive '" + directive + "' must be a valid port number (1 - 65535)");

      directive = "proxy.key";
      proxyKey = config.getString(directive);
      File key = new File(proxyKey);
      if (!key.isFile())
        bailout("configuration directive '" + directive + "' is not a file");
      if (!key.canRead())
        bailout("configuration directive '" + directive + "' is not readable");
      PEMReader in = new PEMReader(new FileReader(key));
      PublicKey publicKey = (PublicKey) in.readObject();
      rsaencrypt.init(Cipher.ENCRYPT_MODE, publicKey);

      directive = "keys.dir";
      keysDir = config.getString(directive);
      File dir = new File(keysDir);
      if (!dir.isDirectory())
        bailout("configuration directive '" + directive + "' is not a directory");
      if (!dir.canRead())
        bailout("configuration directive '" + directive + "' is not readable");
    }
    catch(FileNotFoundException e)
    {
      bailout("unable to read file of directive '" + directive + "'");
    }
    catch(IOException e)
    {
      bailout("while reading proxy public key");
    }
    catch(InvalidKeyException e)
    {
      bailout("invalid key file: " + e.getMessage());
    }
    catch(MissingResourceException e)
    {
      bailout("configuration directive '" + directive + "' is not set");
    }
    catch(NumberFormatException e)
    {
      bailout("configuration directive '" + directive + "' must be numeric");
    }
  }

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

  public void shutdown()
  {
    try
    {
      if (sockin != null)
        sockin.close();
    }
    catch(IOException e)
    {}

    try
    {
      if (sockout != null)
        sockout.close();
    }
    catch(IOException e)
    {}

    try
    {
      if (sock != null && !sock.isClosed())
        sock.close();
    }
    catch(IOException e)
    {}

    try
    {
      if (tPConnection != null)
        tPConnection.join();
    }
    catch(InterruptedException e)
    {}

    try
    {
      if (tInteractive != null)
      {
        tInteractive.interrupt();
        tInteractive.join();
      }
    }
    catch(InterruptedException e)
    {}

    try
    {
      if (stdin != null)
        stdin.close();
    }
    catch(IOException e)
    {}
  }

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

  public void run(String[] args)
  {
    parseArgs(args);
    parseConfig();

    try
    {
      out.println("Connecting to " + proxyHost + ":" + proxyTCPPort + "...");
      sock = new Socket(proxyHost, proxyTCPPort);
      sockin = sock.getInputStream();
      sockout = sock.getOutputStream();
      out.println("Connected...");
    }
    catch(UnknownHostException e)
    {
      bailout("Unable to resolve hostname: " + e.getMessage());
    }
    catch(IOException e)
    {
      bailout("Unable to connect to proxy: " + e.getMessage());
    }

    synchronized(mainLock)
    {
      try
      {
        oin = new Utils.EncObjectInputStream(new BufferedInputStream(sockin));
        oout = new Utils.EncObjectOutputStream(sockout);
        oout.setCipher(rsaencrypt);
      }
      catch(IOException e)
      {
        bailout("Unable to create object stream: " + e.getMessage());
      }

      try
      {
        tPConnection = new Thread(new ProxyConnection(oin, oout,
              sharedFolder, mainLock, interactiveLock));
        tPConnection.start();
      }
      catch(NoSuchMethodException e)
      {
        bailout("Unable to setup remote command handler");
      }

      try
      {
        InputStream stdin = java.nio.channels.Channels.newInputStream(
            new FileInputStream(FileDescriptor.in).getChannel());
        tInteractive = new Thread(new Interactive(stdin, oin, oout,
              mainLock, interactiveLock));
        tInteractive.start();
      }
      catch(NoSuchMethodException e)
      {
        bailout("Unable to setup interactive command handler");
      }

      out.println("Client startup successful!");
      try
      {
        mainLock.wait();
      }
      catch(InterruptedException e)
      {
        /* if we get interrupted -> ignore */
      }

      try
      {
        /* let the threads shutdown */
        Thread.sleep(100);
      }
      catch(InterruptedException e)
      {}
    }

    if (tPConnection != null && !tPConnection.isAlive())
      bailout("Connection to proxy closed unexpected. Terminating...");

    shutdown();
  }

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

  public static void main(String[] args)
  {
    try
    {
      Client cli = new Client();
      cli.run(args);
    }
    catch(Utils.Shutdown e)
    {}
  }
}