/*
 * Copyright (C) 2007-2009 KenD00
 * 
 * This file is part of DumpHD.
 * 
 * DumpHD is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 * 
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 * 
 * You should have received a copy of the GNU General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */
package dumphd.aacs;

import java.security.GeneralSecurityException;
import java.util.concurrent.Semaphore;
import javax.crypto.Cipher;
import javax.crypto.SecretKey;
import javax.crypto.spec.SecretKeySpec;

import dumphd.bdplus.SubTable;
import dumphd.util.PESPack;
import dumphd.util.TSAlignedUnit;
import dumphd.util.TSParserException;

/**
 * This class holds a specific number of threads which are used to decrypt packs of an EVOB / M2TS. This class gets used from the AACSDecrypter.
 * This class is not thread safe, it may be only used by one thread.
 * 
 * Each thread has a working queue in which specific methods of this class put job elements. The workload is automatically balanced between
 * the threads.
 * 
 * TODO: How to handle exceptions that occur in the decrypter threads?
 * TODO: How to handle InterruptedExceptions when adding job elements or waiting for decryption? 
 * TODO: Do work balancing synchronized?
 * 
 * @author KenD00
 */
public final class PackDecrypter {

   /**
    * Array of the used PackDecrypterThreads
    */
   private PackDecrypterThread[] pdts = null;
   /**
    * These are the Thread-Objects created from the PackDecrypterThreads-Objects
    */
   private Thread[] threads = null;
   /**
    * Reference to the supplied packBuffer
    */
   private byte[] packBuffer = null;
   /**
    * Reference to the supplied packGuard
    */
   private Semaphore[] packGuard = null;


   /**
    * Creates a new PackDecrypter using the given number of decrypter threads which each having the given queue length.
    * The decrypter threads are also started.
    * 
    * @param decrypterThreads Number of decrypter threads to create
    * @param queueLength Length in packs of the queue of each decrypter thread
    * @throws GeneralSecurityException An error occured while initializing cryptographic elements
    */
   public PackDecrypter(int decrypterThreads, int queueLength) throws GeneralSecurityException {
      if (decrypterThreads > 0) {
         if (queueLength > 0) {
            pdts = new PackDecrypterThread[decrypterThreads];
            threads = new Thread[decrypterThreads];
            for (int i = 0; i < decrypterThreads; i++) {
               PackDecrypterThread pdt = new PackDecrypterThread(queueLength);
               Thread thread = new Thread(pdt);
               thread.setDaemon(true);
               thread.setPriority(Thread.NORM_PRIORITY + 1);
               thread.start();
               pdts[i] = pdt;
               threads[i] = thread;
            }
         } else throw new IllegalArgumentException("queueLength must be greater than 0");
      } else throw new IllegalArgumentException("decrypterThreads must be greater than 0");
   }

   /**
    * Initializes the object for a new working session.
    * 
    * Waits until all pending jobs are done, then updates the used packBuffer, packGuard and BD+ SubTable.
    * 
    * TODO: Also reset the base address? All actions here don't put anything into the working queue, resetting the base address would.
    * 
    * @param packBuffer The new packBuffer to use
    * @param packGuard The new packGuard to use
    */
   public void init(byte[] packBuffer, Semaphore[] packGuard, SubTable subTable) {
      // TODO: What if we get interrupted during waiting?
      waitForDecryption();
      this.packBuffer = packBuffer;
      this.packGuard = packGuard;
      // Set the PatchIterator for every thread, don't forget to set it to null if the given SubTable is null!
      for (int i = 0; i < pdts.length; i++) {
         if (subTable != null) {
            pdts[i].setPatchIterator(subTable.patchIterator());
         } else {
            pdts[i].setPatchIterator(null);
         }
      }
   }

   /**
    * Updates the base address of the currently used packBuffer.
    * 
    * The base address plus the offset inside the packBuffer is the absolute address of the pack in the source stream.
    * 
    * @param baseAddress The new base address
    */
   public void updateBaseAddress(long baseAddress) {
      try {
         for (int i = 0; i < pdts.length; i++) {
            PackDecrypterThread pdt = pdts[i];
            WorkUnit wu = pdt.acquireWorkUnit();
            wu.type = WorkUnit.UPDATE_BASE_ADDRESS;
            wu.baseAddress = baseAddress;
            pdt.releaseWorkUnit();
         }
      }
      catch (InterruptedException e) {
         // TODO: What to do?
      }
   }
   
   /**
    * Updates the key used for decryption.
    * Adds a new job element to the queue of each decrypter thread with the new key. The key is set after all current jobs in the queue are processed.
    * If a queue is full, this method waites until a place gets free. The key is updated by each decrypter thread itself.
    * 
    * @param npk The new key to use
    * @param encrypt If true, the key will be used in encrypt mode (for BluRay), otherwise in decrypt mode (for HD-DVD)
    */
   public void updateKey(byte[] npk, boolean encrypt) {
      try {
         SecretKey npkKey = new SecretKeySpec(npk, "AES");
         for (int i = 0; i < pdts.length; i++) {
            PackDecrypterThread pdt = pdts[i];
            WorkUnit wu = pdt.acquireWorkUnit();
            if (encrypt) {
               wu.type = WorkUnit.UPDATE_KEY_ENCRYPT;
            }
            else {
               wu.type = WorkUnit.UPDATE_KEY_DECRYPT;
            }
            wu.npkKey = npkKey;
            pdt.releaseWorkUnit();
         }
      }
      catch (InterruptedException e) {
         // TODO: What to do?
      }
   }

   /**
    * Updates the used CPI field.
    * Adds a new job element to the queue of each decrypter thread with the new CPI field. The CPI field is set after all current jobs in the queue are processed.
    * If a queue is full, this method waites until a place gets free. The CPI field is updated by each decrypter thread itself.
    * 
    * @param offset Offset into the packBuffer where the CPI field starts
    */
   public void updateCpi(int offset) {
      try {
         for (int i = 0; i < pdts.length; i++) {
            PackDecrypterThread pdt = pdts[i];
            WorkUnit wu = pdt.acquireWorkUnit();
            wu.type = WorkUnit.UPDATE_CPI;
            // CPI starts at 60, the first 4 bytes (Key Management Information) are not used!
            System.arraycopy(packBuffer, offset + 64, wu.dtk_cpi, 4, 12);
            pdt.releaseWorkUnit();
         }
      }
      catch (InterruptedException e) {
         // TODO: What to do?
      }
   }

   /**
    * Decrypts a pack.
    * Selects the decrypter thread with the lowest workload and adds a new job element to its queue with the given pack offset.
    * If all queues are full, this method waits until a place in the queue of the first decrypter thread gets free.
    * 
    * @param offset Offset into the packBuffer where the pack starts 
    */
   public void decryptPack(int offset) {
      try {
         PackDecrypterThread pdt = getWorker();
         WorkUnit wu = pdt.acquireWorkUnit();
         wu.type = WorkUnit.DECRYPT_PACK;
         wu.packOffset = offset;
         pdt.releaseWorkUnit();
      }
      catch (InterruptedException e) {
         // TODO: What to do?
      }
   }

   public void decryptAlignedUnit(int offset, boolean bdplusOnly) {
      try {
         PackDecrypterThread pdt = getWorker();
         WorkUnit wu = pdt.acquireWorkUnit();
         if (!bdplusOnly) {
            wu.type = WorkUnit.DECRYPT_ALIGNED_UNIT;
         } else {
            wu.type = WorkUnit.REMOVE_BDPLUS;
         }
         wu.packOffset = offset;
         pdt.releaseWorkUnit();
      }
      catch (InterruptedException e) {
         // TODO: What to do?
      }
   }

   /**
    * This method waits until all decrypter threads have finished their jobs.
    */
   public void waitForDecryption() {
      try {
         for (int i = 0; i < pdts.length; i++) {
            pdts[i].waitForIdle();
         }
      }
      catch (InterruptedException e) {
         // TODO: Do we get interrupted?
      }
   }

   /**
    * Terminates all decrypter threads, regardless if they have pending jobs. After this method call this object cannot be used any more.
    */
   public void terminate() {
      if (threads != null) {
         for (int i = 0; i < threads.length; i++) {
            threads[i].interrupt();
            try {
               threads[i].join();
            }
            catch (InterruptedException e) {
               // TODO: Do we get interrupted?
            }
         }
      }
      // Release all used resources
      pdts = null;
      threads = null;
      packBuffer = null;
      packGuard = null;
   }
   
   /**
    * @return The PackDecrypterThread with the lowest workload
    */
   private PackDecrypterThread getWorker() {
      PackDecrypterThread pdt = pdts[0];
      int minCapacity = 0;
      for (int i = 0; i < pdts.length; i++) {
         PackDecrypterThread currentPdt = pdts[i];
         int currentCapacity = currentPdt.getFreeCapacity();
         if (currentCapacity > minCapacity) {
            minCapacity = currentCapacity;
            pdt = currentPdt;
         }
      }
      return pdt;
   }


   /**
    * A PackDecrypterThread processes job elements (WorkUnit objects) put into its queue.
    * 
    * @author KenD00
    */
   private final class PackDecrypterThread implements Runnable {

      /**
       * AES-128 cipher in ECB-Mode
       */
      private Cipher aes_128 = null;
      /**
       * AES-128 cipher in CBC-Mode
       */
      private Cipher aes_128cbcd = null;

      /**
       * Used to parse the decrypted Aligned Unit and set it unencrypted 
       */
      private TSAlignedUnit unit = new TSAlignedUnit();

      /**
       * Base address of the current packBuffer
       */
      private long baseAddress = 0;
      /**
       * DTK_CPI field used for creating of the Content Key
       */
      private byte[] dtk_cpi = new byte[16];
      /**
       * Byte array of the Content Key
       */
      private byte[] ck = new byte[16];
      /**
       * The Content Key
       */
      private SecretKey ckKey = null;
      
      /**
       * The PatchIterator used to iterate over BD+ Patch Entries
       */
      private SubTable.PatchIterator patchIt = null;

      /**
       * Working queue
       */
      private WorkUnit queue[] = null;
      /**
       * The index in the queue where the next WorkUnit should be inserted
       */
      private int putPointer = 0;
      /**
       * The index in the queue where the next WorkUnit should be taken from
       */
      private int getPointer = 0;
      /**
       * Every put operation acquires a permit from this Semaphore. Used to implement waiting if the queue is full
       */
      private Semaphore putGuard = null;
      /**
       * Every get operation acquires a permit from this Semaphore. Used to implement waiting if the queue is empty
       */
      private Semaphore getGuard = new Semaphore(0, false);


      /**
       * Creates a new PackDecrypterThread with the given queue length. Initializes the Cipher objects.
       * 
       * @param queueLength Length in elements of the working queue 
       * @throws GeneralSecurityException An error occurred while initializing cryptographic elements
       */
      public PackDecrypterThread(int queueLength) throws GeneralSecurityException {
         aes_128 = Cipher.getInstance("AES/ECB/NOPADDING");
         aes_128cbcd = Cipher.getInstance("AES/CBC/NOPADDING");
         queue = new WorkUnit[queueLength];
         // Fill the queue with WorkUnits
         // Put and get don't create or delete WorkUnits, the present ones get reused to avoid new / delete overhead
         // The put and get guards prevent overwriting / reusing WorkUnits in a faulty way
         for (int i = 0; i < queue.length; i++) {
            queue[i] = new WorkUnit();
         }
         putGuard = new Semaphore(queueLength, false);
      }
      
      /**
       * Locks a WorkUnit and returns it.
       * 
       * If all WorkUnits are currently locked it waits until one gets released.
       * Every acquired WorkUnit must be released with a call of releaseWorkUnit()!
       * 
       * @return A WorkUnit
       * @throws InterruptedException Interrupted during waiting for a WorkUnit to become available
       */
      public WorkUnit acquireWorkUnit() throws InterruptedException {
         putGuard.acquire();
         WorkUnit wu = queue[putPointer];
         putPointer += 1;
         if (putPointer == queue.length) {
            putPointer = 0;
         }
         return wu;
      }
      
      /**
       * Releases an acquired WorkUnit.
       * 
       * Every acquired WorkUnit must be released. Do not release more WorkUnits than acquired or data loss occurs!
       */
      public void releaseWorkUnit() {
         getGuard.release();
      }
      
      /**
       * @return The available number of free WorkUnits.
       */
      public int getFreeCapacity() {
         return putGuard.availablePermits();
      }
      
      /**
       * Waits until all WorkUnits are processed and this PackDecrypter becomes idle.
       * 
       * @throws InterruptedException Iterrupted during waiting
       */
      public void waitForIdle() throws InterruptedException {
         putGuard.acquire(queue.length);
         putGuard.release(queue.length);
      }
      
      /**
       * Sets the PatchIterator to be used.
       * 
       * WARNING! This method must not be called if there are pending WorkUnits!
       * 
       * TODO: Check if there are no pending WorkUnits?
       * 
       * @param patchIt The PatchIterator to use to patch BD+
       */
      public void setPatchIterator(SubTable.PatchIterator patchIt) {
         this.patchIt = patchIt;
      }

      /* (non-Javadoc)
       * @see java.lang.Runnable#run()
       */
      public void run() {
         WorkUnit wu = null;
         for (;;) {
            try {
               getGuard.acquire();
               wu = queue[getPointer];
               getPointer += 1;
               if (getPointer == queue.length) {
                  getPointer = 0;
               }
               switch (wu.type) {
               case WorkUnit.UPDATE_BASE_ADDRESS:
                  baseAddress = wu.baseAddress;
                  break;
               case WorkUnit.UPDATE_KEY_DECRYPT:
                  try {
                     aes_128.init(Cipher.DECRYPT_MODE, wu.npkKey);
                  }
                  catch (GeneralSecurityException e) {
                     // TODO: What to do?
                  }
                  break;
               case WorkUnit.UPDATE_KEY_ENCRYPT:
                  try {
                     aes_128.init(Cipher.ENCRYPT_MODE, wu.npkKey);
                  }
                  catch (GeneralSecurityException e) {
                     // TODO: What to do?
                  }
                  break;
               case WorkUnit.UPDATE_CPI:
                  // Exchange the arrays to avoid memcopy
                  byte[] temp = dtk_cpi;
                  dtk_cpi = wu.dtk_cpi;
                  wu.dtk_cpi = temp;
                  break;
               case WorkUnit.DECRYPT_PACK:
                  try {
                     // Udate Keyseed
                     // Dtk starts at 84
                     System.arraycopy(packBuffer, wu.packOffset + 84, dtk_cpi, 0, 4);
                     aes_128.doFinal(dtk_cpi, 0, 16, ck);
                     // Do AES-G
                     for (int i = 0; i < 16; i++) {
                        ck[i] = (byte)((byte)ck[i] ^ (byte)dtk_cpi[i]);
                     }
                     ckKey = new SecretKeySpec(ck, "AES");
                     aes_128cbcd.init(Cipher.DECRYPT_MODE, ckKey, AACSDecrypter.cbc_iv_spec);
                     // Encrypted part starts at 128
                     int cryptedOffset = wu.packOffset + 128;
                     aes_128cbcd.doFinal(packBuffer, cryptedOffset, 1920, packBuffer, cryptedOffset);
                  }
                  catch (GeneralSecurityException e) {
                     // TODO: What to do?
                  }
                  // Mark the pack as processed
                  packGuard[wu.packOffset / PESPack.PACK_LENGTH].release();
                  break;
               case WorkUnit.DECRYPT_ALIGNED_UNIT:
                  try {
                     // Perform AACS decryption
                     System.arraycopy(packBuffer, wu.packOffset, dtk_cpi, 0, 16);
                     aes_128.doFinal(dtk_cpi, 0, 16, ck);
                     for (int i = 0; i < 16; i++) {
                        ck[i] = (byte)((byte)ck[i] ^ (byte)dtk_cpi[i]);
                     }
                     ckKey = new SecretKeySpec(ck, "AES");
                     aes_128cbcd.init(Cipher.DECRYPT_MODE, ckKey, AACSDecrypter.cbc_iv_spec);
                     // Encrypted part starts at 16
                     int cryptedOffset = wu.packOffset + 16;
                     aes_128cbcd.doFinal(packBuffer, cryptedOffset, 6128, packBuffer, cryptedOffset);
                     unit.parse(packBuffer, wu.packOffset);
                     unit.setEncrypted(false);
                  }
                  catch (GeneralSecurityException e1) {
                     // TODO: What to do?
                  }
                  catch (TSParserException e2) {
                     // TODO: What to do?
                  }
                  // Fall through to the BD+ removal case
               case WorkUnit.REMOVE_BDPLUS:
                  // Perform BD+ removal, if necessary
                  if (patchIt != null) {
                     // Offset of the patch into current packBuffer
                     long patchOffset = 0;
                     // Check if the Patch Iterator is valid, if not, there are no more patches to apply
                     while (patchIt.isValid()) {
                        // If this test results into true than the patches aren't ordered ascending, this is an error
                        if ((patchOffset = patchIt.getAddress() - baseAddress) < (long)wu.packOffset) {
                           // TODO: Show an error in this situation? Cancel BD+ decryption? Currently the next "valid" patch is searched 
                           patchIt.increment();
                           continue;
                        }
                        // Check if the patch lies in the current Aligned Unit
                        // Need to promote the end of the current Aligned Unit to long because the patchOffset can be more than maxint away
                        if (patchOffset + (long)patchIt.getPatchLength() <= (long)(wu.packOffset + TSAlignedUnit.UNIT_LENGTH)) {
                           // Patch lies in the current Aligned Unit, apply it and increment the iterator
                           // Its safe to cast the offset to an int here because it is never bigger
                           patchIt.getPatch(packBuffer, (int)patchOffset);
                           patchIt.increment();
                        } else {
                           // The patch lies not in the current Aligned Unit. This can have two reasons:
                           // 1. The patch lies in another Aligned Unit. This is OK
                           // 2. The patch exceeds the Aligned Unit (the difference is smaller than the patch length)
                           //    TODO: This is an error situation, currently it is just ignored. This will also raise
                           //          an error that the patches aren't ordered in the next loop
                           break;
                        }
                     }
                  }
                  // Mark the aligned unit as processed
                  packGuard[wu.packOffset / TSAlignedUnit.UNIT_LENGTH].release();
                  break;
               }
               putGuard.release();
            }
            catch (InterruptedException e) {
               return;
            }
         }
      }

   }


   /**
    * An object of this class is a job element in the queue of a decrypter thread.
    * Depending of its type, specific attributes are valid.
    * 
    * @author KenD00
    */
   private final static class WorkUnit {

      /**
       * This WorkUnit represents a pack to be decrypted
       */
      public static final int DECRYPT_PACK = 0;
      /**
       * This WorkUnit represents an aligned unit to be decrypted
       */
      public static final int DECRYPT_ALIGNED_UNIT = 1;
      /**
       * This WorkUnit represents an aligned unit from which only BD+ should be removed
       */
      public static final int REMOVE_BDPLUS = 2;
      /**
       * This WorkUnit represents a Key, Cipher should be initialized in decrypt-mode
       */
      public static final int UPDATE_KEY_DECRYPT = 3;
      /**
       * This WorkUnit represents a Key, Cipher should be initialized in encrypt-mode
       */
      public static final int UPDATE_KEY_ENCRYPT = 4;
      /**
       * This WorkUnit represents a CPI field
       */
      public static final int UPDATE_CPI = 5;
      /**
       * This WorkUnit represents a base address
       */
      public static final int UPDATE_BASE_ADDRESS = 6;

      /**
       * The type of this WorkUnit. This field is always valid.
       */
      public int type = 0;
      /**
       * Base address of the current packBuffer. This field is only valid if this WorkUnit is of the type UPDATE_BASE_ADDRESS
       */
      public long baseAddress = 0;
      /**
       * Offset of a pack. This field is only valid if this WorkUnit is of the type DECRYPT_PACK or DECRYPT_ALIGNED_UNIT or REMOVE_BDPLUS
       */
      public int packOffset = 0;
      /**
       * Reference to a key. This field is only valid if this WorkUnit is of the type UPDATE_KEY_DECRYPT or UPDATE_KEY_ENCRYPT
       */
      public SecretKey npkKey = null;
      /**
       * Reference to a DTK_CPI field. This field is only valid if this WorkUnit is of the type UPDATE_CPI
       */
      public byte[] dtk_cpi = new byte[16];

      /**
       * Default constructor, does nothing
       */
      public WorkUnit(){
         // Nothing
      }
   }

}
