/*
 * A decryptor for encrypted PSP EBOOT binaries.
 * Currently only supports type 2 encryption with the key IDs hardcoded in the KeyStorage class.
 */
using System;
using System.IO;
using System.Collections.Generic;
using System.Security.Cryptography;


class PSPFormatDecryptor
{
	public static readonly uint PSP_MAGIC = 0x5053507E; // "~PSP"
	public static readonly uint PSP_HEADER_SIZE = 336;

	protected class DecryptionHeader
	{
		public byte[] KeyData{ get; set; }
		public byte[] UserKey { get; set; }
		public byte[] Hash { get; set; }
		public byte[] CMACKey { get; set; }
		public uint Tag { get; set; }
		public uint Size { get; set; }
		public uint Offset { get; set; }
	}

	private Stream input;
	private LittleEndianReader reader;

	public static void Main(string[] args)
	{
		if (args.Length == 0) {
			Console.WriteLine("usage: {0} <encrypted> [<decrypted>]", AppDomain.CurrentDomain.FriendlyName);
			return;
		}
		String inpath = args[0];
		String outpath = args.Length < 2 ? "eboot.bin" : args[1];
		using (var instream  = new FileStream(inpath, FileMode.Open))
		using (var outstream = new FileStream(outpath, FileMode.Create)) {
			var decryptor = new PSPFormatDecryptor(instream);
			decryptor.Decrypt(outstream);
		}
	}


	public PSPFormatDecryptor(Stream input)
	{
		this.input = input;
		this.reader = new LittleEndianReader(this.input);
	}

	public void Decrypt(Stream output)
	{
		DecryptionHeader header;
		KeyInfo mainKey;
		byte[] scrambleKey, derivedKey;

		// Parse header and retrieve key.
		this.input.Seek(0, SeekOrigin.Begin);
		header = this.ParseHeader();

		try {
			mainKey = KeyStorage.KEYS[header.Tag];
		} catch (KeyNotFoundException) {
			throw new InvalidOperationException(String.Format("Missing key for tag {0:x} in key storage.", header.Tag));
		}
		this.input.Seek(header.Offset, SeekOrigin.Begin);

		try {
			scrambleKey = KeyStorage.SCRAMBLE_KEYS[mainKey.ScrambleType];
		} catch (KeyNotFoundException) {
			throw new InvalidOperationException(String.Format("Missing descramble key for code 0x{0:x} in key storage.", mainKey.ScrambleType));
		}

		// Get derived key.
		int keyLength = mainKey.Key.Length;
		derivedKey = new byte[9 * keyLength];
		for (int i = 0; i < 9; i++) {
			Array.Copy(mainKey.Key, 0, derivedKey, i * keyLength, keyLength);
			derivedKey[i * keyLength] = (byte)i;
		}
		this.DecryptRound(derivedKey, 0, (uint)(9 * keyLength), derivedKey, 0, scrambleKey);

		// Recreate part of header to decrypt it.
		byte[] origHeader = new byte[0x60];
		int offset = 0;
		offset = this.AppendBuffer(header.KeyData, origHeader, offset);
		offset = this.AppendBuffer(header.Hash,    origHeader, offset);
		offset = this.AppendBuffer(header.UserKey, origHeader, offset);
		offset = this.AppendBuffer(header.CMACKey, origHeader, offset);
		this.DecryptRound(origHeader, 0, 0x60, origHeader, 0, scrambleKey);

		// Now decrypt key from decrypted header.
		byte[] userKey = new byte[0x10];
		Array.Copy(origHeader, header.KeyData.Length + header.Hash.Length, userKey, 0, userKey.Length);
		for (int i = 0; i < userKey.Length; i++)
			userKey[i] ^= derivedKey[keyLength + i];
		this.DecryptRound(userKey, 0, (uint)userKey.Length, userKey, 0, scrambleKey);
		for (int i = 0; i < userKey.Length; i++)
			userKey[i] ^= derivedKey[0x40 + keyLength + i];
		this.DecryptRound(userKey, 0, (uint)userKey.Length, userKey, 0, KeyStorage.MASTER_KEY);

		// And decrypt the data now that we have the key.
		this.input.Seek(header.Offset, SeekOrigin.Begin);
		this.DecryptRound(this.input, header.Size, output, userKey);
	}

	protected int AppendBuffer(byte[] buf, byte[] output, int offset)
	{
		Array.Copy(buf, 0, output, offset, buf.Length);
		return offset + buf.Length;
	}

	protected DecryptionHeader ParseHeader()
	{
		DecryptionHeader header = new DecryptionHeader();
		long pos = this.input.Position;

		uint magic = this.reader.ReadU32();
		if (magic != PSP_MAGIC)
			throw new InvalidOperationException("Input file is not a PSP format file (missing magic).");

		// Attributes.
		this.reader.ReadU32();
		// Version (low and high).
		this.reader.ReadU16();
		// Name.
		for (int i = 0; i < 28; i++)
			this.reader.ReadU8();
		// Version.
		this.reader.ReadU8();
		// Segment count.
		this.reader.ReadU8();
		// ELF size.
		header.Size = this.reader.ReadU32();
		// PSP size.
		this.reader.ReadU32();
		// Boot entry.
		this.reader.ReadU32();
		// Module information offset.
		this.reader.ReadU32();
		// BSS segment size.
		this.reader.ReadU32();
		// Segment alignments.
		for (int i = 0; i < 4; i++)
			this.reader.ReadU16();
		// Segment offsets;
		for (int i = 0; i < 4; i++)
			this.reader.ReadU32();
		// Segment sizes.
		for (int i = 0; i < 4; i++)
			this.reader.ReadU32();
		// Reserved.
		for (int i = 0; i < 5; i++)
			this.reader.ReadU32();
		// SDK version.
		this.reader.ReadU32();
		// Mode.
		this.reader.ReadU8();
		// Padding.
		this.reader.ReadU8();
		// Size overlaps.
		this.reader.ReadU16();
		// AES user key.
		header.UserKey = new byte[16];
		this.reader.ReadBytes(header.UserKey.Length, header.UserKey);
		// CMAC key.
		header.CMACKey = new byte[16];
		this.reader.ReadBytes(header.CMACKey.Length, header.CMACKey);
		// CMAS header hash.
		for (int i = 0; i < 16; i++)
			this.reader.ReadU8();
		// CMAC unknown hash.
		for (int i = 0; i < 16; i++)
			this.reader.ReadU8();
		// CMAC data hash.
		for (int i = 0; i < 16; i++)
			this.reader.ReadU8();
		// Tag.
		header.Tag = this.reader.ReadU32();
		// Signature.
		for (int i = 0; i < 88; i++)
			this.reader.ReadU8();
		// SHA-1 hash.
		header.Hash = new byte[20];
		this.reader.ReadBytes(header.Hash.Length, header.Hash);
		// Key data.
		header.KeyData = new byte[16];
		this.reader.ReadBytes(header.KeyData.Length, header.KeyData);

		// End.
		header.Offset = (uint)this.input.Position;
		if (header.Offset != PSP_HEADER_SIZE)
			throw new InvalidOperationException("Input file is not a PSP format file (malformed header).");

		return header;
	}

	protected void DecryptRound(byte[] buf, uint offset, uint size, byte[] output, uint ooffset, byte[] key)
	{
		Stream instream = new MemoryStream(buf, (int)offset, (int)size);
		Stream outstream = new MemoryStream(output, (int)ooffset, (int)(output.Length - ooffset));
		DecryptRound(instream, size, outstream, key);
	}

	protected void DecryptRound(Stream input, uint size, Stream output, byte[] key)
	{
		using (var algorithm = new RijndaelManaged()) {
			algorithm.BlockSize = 128;
			algorithm.KeySize = 128;
			algorithm.Mode = CipherMode.CBC;
			algorithm.Padding = PaddingMode.None;
			algorithm.Key = key;
			algorithm.IV = new byte[16];
			using (var decryptor = algorithm.CreateDecryptor()) {
				byte[] inbuf = new byte[decryptor.InputBlockSize];
				byte[] outbuf = new byte[decryptor.OutputBlockSize];

				for (int i = 0, n = 0; i < size; i += n) {
					n = input.Read(inbuf, 0, (int)Math.Min(inbuf.Length, size - i));
					if (n == 0)
						break;
					int d = decryptor.TransformBlock(inbuf, 0, n, outbuf, 0);
					output.Write(outbuf, 0, d);
				}
			}
		}
	}
}


class KeyInfo
{
	public byte[] Key { get; set; }
	public int Type { get; set; }
	public uint ScrambleType { get; set; }

	public KeyInfo(byte[] key, int type, uint scrambleType)
	{
		this.Key = key;
		this.Type = type;
		this.ScrambleType = scrambleType;
	}
}

class KeyStorage
{
	public static readonly Dictionary<uint, KeyInfo> KEYS = new Dictionary<uint, KeyInfo>() {
		{ 0xD91613F0, new KeyInfo(new byte[] { 0xEB, 0xFF, 0x40, 0xD8, 0xB4, 0x1A, 0xE1, 0x66, 0x91, 0x3B, 0x8F, 0x64, 0xB6, 0xFC, 0xB7, 0x12 }, 0x2, 0x5D) }
	};
	public static readonly Dictionary<uint, byte[]> SCRAMBLE_KEYS = new Dictionary<uint, byte[]>() {
		{ 0x5D, new byte[] { 0x11, 0x5A, 0x5D, 0x20, 0xD5, 0x3A, 0x8D, 0xD3, 0x9C, 0xC5, 0xAF, 0x41, 0x0F, 0x0F, 0x18, 0x6F } }
	};
	public static readonly byte[] MASTER_KEY = new byte[] {
		0x98, 0xC9, 0x40, 0x97, 0x5C, 0x1D, 0x10, 0xE8, 0x7F, 0xE6, 0x0E, 0xA3, 0xFD, 0x03, 0xA8, 0xBA 
	};
}


class LittleEndianReader
{
	private Stream input;

	public LittleEndianReader(Stream input)
	{
		this.input = input;
	}

	public uint ReadU8() { return (uint)this.input.ReadByte(); }
	public uint ReadU16() { return this.ReadU8() | (this.ReadU8() << 8); }
	public uint ReadU32() { return this.ReadU16() | (this.ReadU16() << 16); }
	public void ReadBytes(int length, byte[] buf) { this.input.Read(buf, 0, length); }
}