package com.mpaa.decss;

import java.io.*;


/**
 * Title:        TitleKeyCracker.java
 * Description:  This class contains a cryptographic attack on the encrypted content of a DVD (Content Scrambling System) to obtain the
 *               title key.
 * @author DiatriBe
 * @version 1.0
 */

public class TitleKeyCracker
{
	File inputVobFile;
	LinearFeedbackShiftRegisterPair feedbackRegisters;

	/**
	 * This array contains the data pattern of the special block that we are looking for in order to crack it (we can
	 * crack these blocks because we know the expected plain text output of certain encrypted bytes in this block).
	 *
	 * Format of the data in this table is 'sector offset', 'desired value'
	 */
	private static int[] crackableBlockPattern =
	{
		0x00, 0x00,
		0x01, 0x00,
		0x02, 0x01,
		0x03, 0xBA,
		0x0E, 0x00,
		0x0F, 0x00,
		0x10, 0x01
	};

	/**
	 * This is the plain version of the enctrypted bytes that we know the value of in our 'special' block.  Note that
	 * byte 4 and 5 are set to 0 here but actually have a value which is specific to the value of the block in which
	 * they are present.  The value for these bytes will be set on the fly based on the contents of the 'special' block.
	 */
	private static int plainText[] = { 0x00, 0x00, 0x01, 0xBE, 0x00, 0x00, 0xFF };


	/**
	 * Constructor that takes an input file that will be searched to find the title key.
	 *
	 * @param inputVobFile file for which to attempt to apply the title key crack.
	 */
	public TitleKeyCracker(File inputVobFile)
	{
		this.inputVobFile = inputVobFile;
	}

	/**
	 * This method will attempt to crack the CSS encryption by searching for a 'special' block which contains encrypted data
	 * for which we know the plain-text equivalent.
	 *
	 * @return boolean - indicates whether the title key crack succeeded for this file.
	 */
	public boolean attackKey()
	{
		boolean keyFound = false;
		boolean foundProperSector = false;
		boolean foundLikelySector = false;

		System.out.println("Scanning for key in file '" + inputVobFile.getName() + "'");

		try
		{
			DataInputStream reader = new DataInputStream(new BufferedInputStream(new FileInputStream(inputVobFile), 0x400000));        // Allocate a 4 MB input buffer

			long fileSize = inputVobFile.length();
			long bytesRead = 0;

			int[] sector = new int[2048];


			for (int block = 0; (bytesRead < fileSize) && !keyFound; block++)
			{
				for (int index = 0; index < 2048; index++)
					sector[index] = reader.readUnsignedByte();

				bytesRead += 2048;

				if ((block % 1000) == 0)
					System.out.println("Scanning block " + block);

				if ((sector[0x14] & 0x10) != 0)
				{
					foundLikelySector = true;

					boolean patternMatched = true;
					for (int index = 0; (index < crackableBlockPattern.length) && patternMatched; index += 2)
					{
						if (sector[crackableBlockPattern[index]] != crackableBlockPattern[index + 1])
							patternMatched = false;
					}
					if (patternMatched)
					{
						int offset = 0x14 + (sector[0x12] << 8) + sector[0x13];
						if ((offset >= 0x80) && (offset <= 0x7F9))
						{
							foundProperSector = true;
							int left = 0x800 - offset - 6;
							plainText[4] = (left >> 8) & 0xFF;
							plainText[5] = left & 0xFF;
							int[] cipherText = { sector[offset], sector[offset + 1], sector[offset + 2], sector[offset + 3], sector[offset + 4], sector[offset + 5], sector[offset + 6] };
							int count = findKeys(cipherText, plainText, offset);
							if (count > 0)     // count = the number of likely keys found based on the attack of this sector
							{
								if (count == 1)     // Found only one key, looks good!  So we can now break out of the loop.
								{
									int[] salt = { sector[0x54], sector[0x55], sector[0x56], sector[0x57], sector[0x58] };
									feedbackRegisters.salt(salt);       // Now just cancel out the sector salting and what remains should be the title key
									keyFound = true;
								}
								else
								{
									feedbackRegisters = null;       // Clear out the bad value for the feedback registers (since we are ignoring them)
									System.out.println("Block " + block + " reported " + count + " possible keys so we are skipping it");
								}
							}
						}
					}
				}
			}
			reader.close();

			System.out.println("The cryptographic attack is done");
			if (!keyFound)
			{
				System.out.print("Error: Unable to attack key from input file (");

				if (foundProperSector)
					System.out.print("the algorithm failed!");
				else if (foundLikelySector)
					System.out.print("there were no vulnerable blocks");
				else
					System.out.print("there were no encrypted blocks");
				System.out.println(")");
			}
		}
		catch (FileNotFoundException ex)
		{
			System.out.println("The imput file '" + inputVobFile.getName() + "' could not be found.");
		}
		catch (Exception ex)
		{
			System.out.println("There was an error encountered while trying to find the title key in file '" + inputVobFile.getName() + "'");
		}
		return keyFound;
	}

	/**
	 * Will return the title key that is found (if one was not found, will return a null).
	 *
	 * @return TitleKey - the title key that was found by applying the cipher crack.  Will return null if no key could be found.
	 */
	public TitleKey getTitleKey()
	{
		if (feedbackRegisters != null)
			return feedbackRegisters.getKey();
		else
			return null;
	}

	/**
	 * The following code is a (modified) implementation of an algorithm originally described by Frank Stevenson.
	 * A synopsis of this description is:
	 *
	 * The CSS algorithm is fatally flawed. A divide and conquer attack is possible by guessing the 17 unknown bits of
	 * LFSR17. LFSR17 is then clocked 4 times, and the known keystream bytes are then used to reconstruct the state of
	 * LFSR25. The whole cipher is then clocked another 2-6 times to validate the guess. If the LFSR17 guess is correct,
	 * both shift registers are clocked backwards to retrieve the initial state (which is the title key).
	 *
	 * @param cipherText the cipher-text data that will be used for the crack.
	 * @param plainText the plain-text data that will be used for the crack.
	 * @param offset the number of bytes that the cipher text is offset from the beginning of the data block.
	 *
	 * @return int - count of the number of potential keys found for this data block.  If more than one key was found,
	 *               this block should be ignored since no unique key usually = no correct key.
	 */
	private int findKeys(int[] cipherText, int[] plainText, int offset)
	{
		LinearFeedbackShiftRegisterPair testFeedbackRegisters = new LinearFeedbackShiftRegisterPair();

		if ((cipherText.length < 7) || (plainText.length < 7))
		{
			System.out.println("Not enough bytes for the find title key attack!");
			return 0;
		}

		// By reversing what is done to calculate the plain text from the cipher text, we can calculate what the cipher stream output needs to be (this will correspond to the calculation of lfsr17 + lfsr25 + carry)
		int[] cipherStreamOutput = new int[7];
		for (int index = 0; index < 7; index++)
			cipherStreamOutput[index] = SubstitutionTable.getAt(cipherText[index]) ^ plainText[index];

		int count = 0;      // Number of keys found based on attack of this sector

		// Go through every possible initial state for LFSR0 until a match for the cipherStreamOutput is found (there are 2^18 permutations)
		for (int permutation = 0; permutation < 0x40000; permutation++)
		{
			testFeedbackRegisters.setLfsr17Register(permutation >> 1);        // Use the LSB of permutation as the carry bit (so block it out here)
			testFeedbackRegisters.setLfsr25Register(0);
			int carry = permutation & 0x01;                   // Now set the LSB of permutation as the carry bit

			// By clocking lfsr17 forward 4 iterations and for each iteration taking its output and reversing the conversion
			// that is done to create the cipher stream, we can extract what the lfsr27 output would have to be for the
			// cipher stream output to look like it does.  Once we have done this 4 times (with 4 bytes) lfsr27 will be fully loaded.
			for (int index = 0; index < 4; index++)
			{
				testFeedbackRegisters.clockLfsr17OneByte();
				int lfsr27Output = cipherStreamOutput[index] - testFeedbackRegisters.getLfsr17Output(true) - carry;
				testFeedbackRegisters.setLfsr25Register((testFeedbackRegisters.getLfsr25Register() >> 8) | ((lfsr27Output & 0xFF) << 17));
				carry = (lfsr27Output >> 8) & 0x01;
			}

			// Now that we have a fully loaded lfsr17 and lfsr27, we can verify that they produce the proper output by clocking them
			// 3 times and comparing their output to the known cipher stream.
			boolean foundPossible = true;
			for (int index = 4; (index < 7) && foundPossible; index++)
			{
				testFeedbackRegisters.clockOneByte();
				int testCipherStreamOutput = testFeedbackRegisters.getLfsr17Output(true) + testFeedbackRegisters.getLfsr25Output(false) + carry;
				if ((testCipherStreamOutput & 0xFF) != cipherStreamOutput[index])
					foundPossible = false;

				carry = testCipherStreamOutput >> 8;
			}
			if (foundPossible)
			{
				// Now clock back the registers to the start state and verify that we have the carry bit set right
				for (int index = 5; index >= 0; index--)
					testFeedbackRegisters.clockBackwardOneByte();

				int testCipherStreamOutput = testFeedbackRegisters.getLfsr17Output(true) + testFeedbackRegisters.getLfsr25Output(false) + (permutation & 0x01);

				if ((testCipherStreamOutput & 0xFF) == cipherStreamOutput[0])
				{
					// Clock backwards to the beginning of the sector (should bring the shift registers to their initial state)
					for (int index = offset; index >= 128; index--)
						testFeedbackRegisters.clockBackwardOneByte();

					// Test if we found a title key, based on the check that the feedback register bits that should be set (to avoid null cycling) at initialization actually are set!
					if (((testFeedbackRegisters.getLfsr17Register() & 0x100) != 0) && ((testFeedbackRegisters.getLfsr25Register() & 0x200000) != 0))
					{
						try
						{
							feedbackRegisters = (LinearFeedbackShiftRegisterPair) testFeedbackRegisters.clone();
						}
						catch (CloneNotSupportedException ex)       // Man, I hope this never happens!
						{
							ex.printStackTrace();
						}
						count++;
					}
				}
			}
		}
		return count;
	}

	public static void main(String[] args)
	{
		if (args.length != 1)
		{
			System.out.println("Wrong number of arguments!");
			return;
		}

		TitleKeyCracker titleKeyCracker = new TitleKeyCracker(new File(args[0]));

		titleKeyCracker.attackKey();

		TitleKey titleKey = titleKeyCracker.getTitleKey();

		System.out.println("Key = " + titleKey.toString());
	}
}