/*--
  This file is a part of ZetaGrid, a simple and secure Grid Computing
  kernel.

  Copyright (c) 2001-2004 Sebastian Wedeniwski.  All rights reserved.

  Use in source and binary forms, with or without modification,
  are permitted provided that the following conditions are met:

  1. The source code must retain the above copyright
     notice, this list of conditions and the following disclaimer.

  2. The origin of this software must not be misrepresented; you must 
     not claim that you wrote the original software.  If you plan to
     use this software in a product, please contact the author.

  3. Altered source versions must be plainly marked as such, and must
     not be misrepresented as being the original software. The author
     must be informed about these changes.

  4. The name of the author may not be used to endorse or promote 
     products derived from this software without specific prior written 
     permission.

  THIS SOFTWARE IS PROVIDED BY THE AUTHOR ``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 AUTHOR 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.

  Version 1.9.0, February 8, 2004

  This program is based on the work of:
     S. Wedeniwski
--*/

package zeta.tool;

import java.io.BufferedReader;
import java.io.BufferedWriter;
import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.File;
import java.io.FileInputStream;
import java.io.FileOutputStream;
import java.io.FileReader;
import java.io.FileWriter;
import java.io.InputStream;
import java.io.IOException;
import java.io.StringReader;
import java.io.StringWriter;
import java.io.OutputStream;
import java.math.BigInteger;
import java.sql.Connection;
import java.sql.DriverManager;
import java.sql.PreparedStatement;
import java.sql.ResultSet;
import java.sql.SQLException;
import java.sql.Statement;

import zeta.crypto.IONumber;
import zeta.crypto.KeyManager;
import zeta.crypto.Signature;
import zeta.util.Base64;
import zeta.util.DatabaseUtils;
import zeta.util.StreamUtils;
import zeta.util.StringUtils;

public class NewVersion {
  public static void main(String[] args) {
    try {
      if (args.length == 6 && args[0].equals("uec")) {
        NewVersion newVersion = new NewVersion(args[4], args[5]);
        newVersion.updateEncryptionClass(Integer.parseInt(args[1]), args[2], Integer.parseInt(args[3]));
        newVersion.close();
      } else if (args.length >= 6 && args.length <= 9 && args[0].length() == 1) {
        char c = args[0].charAt(0);
        String user = args[args.length-2];
        String password = args[args.length-1];
        String os = "";
        String arch = "";
        int taskId = 0;
        int randomize = Integer.parseInt(args[args.length-3]);
        int i = 1;
        if (args.length >= 8) {
          os = args[args.length-5];
          arch = args[args.length-4];
        }
        if (args.length == 7 || args.length == 9) {
          taskId = Integer.parseInt(args[1]);
          i = 2;
        }
        NewVersion newVersion = new NewVersion(user, password);
        if (c == 'i') {
          newVersion.insertNewVersion(taskId, args[i], args[i+1], os, arch, randomize, null);    // ToDo: configure key class name
        } else if (c == 'u') {
          newVersion.updateNewVersion(taskId, args[i], args[i+1], os, arch, randomize, null);
        } else if (c == 'd') {
          newVersion.deleteVersion(taskId, args[i], args[i+1], os, arch, randomize);
        } else if (c == 's') {
          newVersion.simpleUpdate(taskId, args[i], args[i+1], os, arch, (randomize != 0));
        }
        newVersion.close();
      } else {
        System.err.println("USAGE: java zeta.tool.NewVersion <{i,u,d,s}> [<taskId>] <version> <program> [<os> <arch>] {<randomize>,<compress>} <user> <password>\n"
                         + "       java zeta.tool.NewVersion uec <taskId> <encryption class filename> <randomize> <user> <password>\n"
                         + "       i   - insert a new program and update the signatures\n"
                         + "       u   - update the program and update the signatures\n"
                         + "       d   - delete the program and update the signatures\n"
                         + "       s   - update the program\n"
                         + "       uec - update the encryption class");
      }
    } catch (Exception e) {
      e.printStackTrace();
    }
  }

  private NewVersion(String user, String password) throws Exception {
    Class.forName("COM.ibm.db2.jdbc.app.DB2Driver").newInstance();
    connection = DriverManager.getConnection("jdbc:db2:zeta", user, password);
  }

  private void close() {
    DatabaseUtils.close(connection);
  }

  private static String getFilename(String program, String osName, String arch) {
    if (osName.length() == 0) {
      return program;
    }
    String name = osName + '/' + arch + '/' + program;
    File file = new File(name);
    return (file.exists())? name : (osName + '/' + program);
  }

  /**
   *  Returns (version, all signatures)
  **/
  private String[] updateSignature(boolean delete, String program, String osName, String arch, int randomize, String keyClassname) throws Exception {
    Signature signature = new Signature(KeyManager.getKey(keyClassname));
    InputStream in = null;
    InputStream privateKeyFile  = null;
    Statement stmt = null;
    try {
      ByteArrayOutputStream out = new ByteArrayOutputStream(16*1024);
      if (!delete) {
        in = new FileInputStream(getFilename(program, osName, arch));
      }
      privateKeyFile = new FileInputStream("private_key.txt");    // ToDo: must be a different private key for keyClassname
      BigInteger privateKey = IONumber.read(privateKeyFile);
      if (privateKey == null) {
        throw new IOException("Missing private key file 'private_key.txt'");
      }
      if (!delete) {
        signature.generate(randomize, privateKey, in, out);
      }
      String newSignature = out.toString("UTF-8");

      out.reset();
      stmt = connection.createStatement();
      ResultSet rs = stmt.executeQuery("SELECT version,program FROM zeta.program WHERE task_id=0 AND name='signature.txt'");
      String version = "0100";
      if (!rs.next()) {
        rs.close();
        rs = stmt.executeQuery("SELECT COUNT(*) FROM zeta.program");
        if (rs.next() && rs.getInt(1) > 0) {
          throw new SQLException("'signature.txt' not in db!");
        }
        rs.close();
      } else {
        version = StringUtils.format(false, '0', 4, Integer.toString((Integer.parseInt(rs.getString(1).trim())+1)));
        StreamUtils.writeData(rs.getBinaryStream(2), out, true, true);
        rs.close();
      }
      StringWriter sWriter = new StringWriter();
      BufferedReader reader = new BufferedReader(new StringReader(out.toString("UTF-8")));
      BufferedWriter writer = new BufferedWriter(sWriter);
      while (true) {
        String filename = reader.readLine();
        if (filename == null) {
          break;
        }
        boolean found = false;
        int idx = filename.indexOf(',');
        int idx2 = filename.indexOf(',', idx+1);
        if (idx2+1 == filename.length()) {
          filename = filename.substring(0, idx2);
          idx2 = filename.indexOf(',', idx+1);
        }
        if (idx < idx2 && idx > 0 && idx2+1 < filename.length()) {
          found = (filename.substring(idx+1, idx2).equalsIgnoreCase(osName) && filename.substring(0, idx).equalsIgnoreCase(program)
                   && filename.substring(idx2+1).equalsIgnoreCase(arch));
        } else if ((idx == -1 || idx == idx2) && filename.equalsIgnoreCase(program)) {
          found = true;
        }
        if (found) {
          reader.readLine();
        } else {
          String s = reader.readLine();
          if (s != null) {
            writer.write(filename);
            writer.newLine();
            writer.write(s);
            writer.newLine();
          }
        }
      }
      reader.close();
      if (!delete) {
        writer.write(program);
        if (osName.length() > 0 && arch.length() > 0) {
          writer.write(",");
          writer.write(osName);
          writer.write(",");
          writer.write(arch);
        }
        writer.newLine();
        writer.write(newSignature);
      }
      writer.newLine();
      writer.close();
      String allSignatures = sWriter.toString();
      out.reset();
      signature = new Signature(KeyManager.getKey(null));
      signature.generate(randomize, privateKey, new ByteArrayInputStream(allSignatures.getBytes("UTF-8")), out);
      return new String[] { version,  allSignatures + out.toString("UTF-8") };
    } finally {
      DatabaseUtils.close(stmt);
      StreamUtils.close(in);
      StreamUtils.close(privateKeyFile);
    }
  }

  private void insertNewVersion(int taskId, String version, String program, String os, String arch, int randomize, String keyClassname) throws Exception {
    PreparedStatement stmt = null;
    try {
      String[] res = updateSignature(false, program, os, arch, randomize, keyClassname);
      String versionSig = res[0];
      byte[] signature = res[1].getBytes("UTF-8");
      // ToDo: INSERT key_classname into database
      String s = "INSERT INTO zeta.program (task_id,name,os_name,os_version,os_arch,version,compressed_YN,program,last_update) VALUES ("
                 + taskId + ",'" + program + "','" + os + "','','" + arch + "','" + version + "','Y',?,CURRENT TIMESTAMP)";
      System.out.println(s);
      stmt = connection.prepareStatement(s);
      stmt.setBytes(1, StreamUtils.getFile(getFilename(program, os, arch), true, true));
      System.out.println("exec");
      stmt.executeUpdate();
      stmt.close();
      stmt = connection.prepareStatement("SELECT COUNT(*) FROM zeta.program WHERE name='signature.txt'");
      ResultSet rs = stmt.executeQuery();
      if (rs.next() && rs.getInt(1) == 0) {
        rs.close();
        System.out.println("insert signature");
        stmt = connection.prepareStatement("INSERT INTO zeta.program (task_id,name,os_name,os_version,os_arch,version,compressed_YN,program,last_update) VALUES (0,'signature.txt','','','','"
                                            + versionSig + "','N',?,CURRENT TIMESTAMP)");
      } else {
        rs.close();
        System.out.println("update signature");
        stmt = connection.prepareStatement("UPDATE zeta.program SET (version,compressed_YN,program,last_update)=('" + versionSig
                     + "','N',?,CURRENT TIMESTAMP) WHERE task_id=0 AND name='signature.txt' AND os_name='' AND os_version='' AND os_arch=''");
      }
      stmt.setBytes(1, signature);
      stmt.executeUpdate();
      System.out.println("ready");
      System.out.println(res[1]);
    } finally {
      DatabaseUtils.close(stmt);
    }
  }

  private void updateNewVersion(int taskId, String version, String program, String os, String arch, int randomize, String keyClassname) throws Exception {
    PreparedStatement stmt = null;
    try {
      String[] res = updateSignature(false, program, os, arch, randomize, keyClassname);
      String versionSig = res[0];
      byte[] signature = res[1].getBytes("UTF-8");
      // ToDo: UPDATE key_classname into database
      String s = "UPDATE zeta.program SET (version,compressed_YN,program,last_update)=(?,'Y',?,CURRENT TIMESTAMP) WHERE name=? AND task_id="
                 + taskId + " AND os_name='" + os + "' AND os_version='' AND os_arch='" + arch + '\'';
      System.out.println(s);
      stmt = connection.prepareStatement(s);
      stmt.setString(1, version);
      stmt.setBytes(2, StreamUtils.getFile(getFilename(program, os, arch), true, true));
      stmt.setString(3, program);
      System.out.println("exec");
      int result = stmt.executeUpdate();
      System.out.println("result=" + result);
      stmt.close();
      stmt = null;
      if (result != 1) {
        System.out.println("ERROR!");
      } else {
        System.out.println("update signature");
        stmt = connection.prepareStatement("UPDATE zeta.program SET (version,compressed_YN,program,last_update)=('" + versionSig
                     + "','N',?,CURRENT TIMESTAMP) WHERE task_id=0 AND name='signature.txt' AND os_name='' AND os_version='' AND os_arch=''");
        stmt.setBytes(1, signature);
        stmt.executeUpdate();
        System.out.println("ready");
        System.out.println(res[1]);
      }
    } finally {
      DatabaseUtils.close(stmt);
    }
  }

  private void deleteVersion(int taskId, String version, String program, String os, String arch, int randomize) throws Exception {
    PreparedStatement stmt = null;
    try {
      String[] res = updateSignature(true, program, os, arch, randomize, null);
      String versionSig = res[0];
      byte[] signature = res[1].getBytes("UTF-8");
      String s = "DELETE FROM zeta.program WHERE name=? AND task_id=" + taskId + " AND os_name='" + os + "' AND os_version='' AND os_arch='" + arch + '\'';
      System.out.println(s);
      stmt = connection.prepareStatement(s);
      stmt.setString(1, program);
      System.out.println("exec");
      int result = stmt.executeUpdate();
      System.out.println("result=" + result);
      stmt.close();
      stmt = null;
      if (result != 1) {
        System.out.println("Warning: program is not defined!");
      }
      System.out.println("update signature");
      /*ByteArrayOutputStream out = new ByteArrayOutputStream();
      StreamUtils.writeData(new FileInputStream("backup/zeta.program.5.zip"), out, true, true);
      byte[] signature = out.toByteArray();*/
      stmt = connection.prepareStatement("UPDATE zeta.program SET (version,compressed_YN,program,last_update)=('" + versionSig
                   + "','N',?,CURRENT TIMESTAMP) WHERE task_id=0 AND name='signature.txt' AND os_name='' AND os_version='' AND os_arch=''");
      stmt.setBytes(1, signature);
      stmt.executeUpdate();
      System.out.println("ready");
      System.out.println(res[1]);
    } finally {
      DatabaseUtils.close(stmt);
    }
  }

  private void simpleUpdate(int taskId, String version, String program, String os, String arch, boolean compress) throws Exception {
    PreparedStatement stmt = null;
    try {
      String s = "UPDATE zeta.program SET (version,compressed_YN,program)=(?,'" + ((compress)? 'Y' : 'N')
                 + "',?) WHERE name=? AND task_id=" + taskId
                 + " AND os_name='" + os + "' AND os_version='' AND os_arch='" + arch + '\'';
      System.out.println(s);
      stmt = connection.prepareStatement(s);
      stmt.setString(1, version);
      stmt.setBytes(2, StreamUtils.getFile(getFilename(program, os, arch), compress, true));
      stmt.setString(3, program);
      System.out.println("exec");
      int result = stmt.executeUpdate();
      System.out.println("result=" + result);
      if (result != 1) {
        System.out.println("ERROR!");
      }
    } finally {
      DatabaseUtils.close(stmt);
    }
  }

  private void updateEncryptionClass(int taskId, String classFileName, int randomize) throws Exception {
    InputStream in = null;
    InputStream privateKeyFile = null;
    PreparedStatement stmt = null;
    Signature signature = new Signature(KeyManager.getKey(null));
    try {
      ByteArrayOutputStream out = new ByteArrayOutputStream(4*1024);
      in = new FileInputStream(classFileName);
      StreamUtils.writeData(in, out, true, true);
      in = null;
      byte[] encryptionClass = out.toByteArray();
      privateKeyFile = new FileInputStream("private_key.txt");    // ToDo: must be a different private key for keyClassname
      BigInteger privateKey = IONumber.read(privateKeyFile);
      if (privateKey == null) {
        throw new IOException("Missing private key file 'private_key.txt'");
      }
      out.reset();
      signature.generate(randomize, privateKey, new ByteArrayInputStream(encryptionClass), out);
      String s = "UPDATE zeta.task SET (encryption_class,encryption_signature)=(?,?) WHERE id=" + taskId;
      System.out.println(s);
      stmt = connection.prepareStatement(s);
      stmt.setString(1, Base64.encode(encryptionClass));
      stmt.setString(2, out.toString("UTF-8"));
      System.out.println("exec");
      int result = stmt.executeUpdate();
      if (result != 1) {
        System.out.println("ERROR! result=" + result);
      }
    } finally {
      DatabaseUtils.close(stmt);
      StreamUtils.close(in);
    }
  }

  private Connection connection = null;
}