#region PDFsharp - A .NET library for processing PDF
//
// Authors:
//   Stefan Lange (mailto:Stefan.Lange@pdfsharp.com)
//
// Copyright (c) 2005-2008 empira Software GmbH, Cologne (Germany)
//
// http://www.pdfsharp.com
// http://sourceforge.net/projects/pdfsharp
//
// Permission is hereby granted, free of charge, to any person obtaining a
// copy of this software and associated documentation files (the "Software"),
// to deal in the Software without restriction, including without limitation
// the rights to use, copy, modify, merge, publish, distribute, sublicense,
// and/or sell copies of the Software, and to permit persons to whom the
// Software is furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included
// in all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
// THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
// FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER 
// DEALINGS IN THE SOFTWARE.
#endregion

using System;
using System.Diagnostics;
using System.Collections;
using PdfSharp.Drawing;
using PdfSharp.Internal;
using PdfSharp.Pdf;
using PdfSharp.Pdf.IO;
using PdfSharp.Pdf.Advanced;
using PdfSharp.Pdf.Internal;
using System.Security.Cryptography;

namespace PdfSharp.Pdf.Security
{
  /// <summary>
  /// Represents the standard PDF security handler.
  /// </summary>
  public sealed class PdfStandardSecurityHandler : PdfSecurityHandler
  {
    internal PdfStandardSecurityHandler(PdfDocument document)
      : base(document)
    {
    }

    internal PdfStandardSecurityHandler(PdfDictionary dict)
      : base(dict)
    {
    }

    /// <summary>
    /// Sets the user password of the document. Setting a password automatically sets the
    /// PdfDocumentSecurityLevel to PdfDocumentSecurityLevel.Encrypted128Bit if its current
    /// value is PdfDocumentSecurityLevel.None.
    /// </summary>
    public string UserPassword
    {
      set
      {
        if (this.document.securitySettings.DocumentSecurityLevel == PdfDocumentSecurityLevel.None)
          this.document.securitySettings.DocumentSecurityLevel = PdfDocumentSecurityLevel.Encrypted128Bit;
        this.userPassword = value;
      }
    }
    internal string userPassword;

    /// <summary>
    /// Sets the owner password of the document. Setting a password automatically sets the
    /// PdfDocumentSecurityLevel to PdfDocumentSecurityLevel.Encrypted128Bit if its current
    /// value is PdfDocumentSecurityLevel.None.
    /// </summary>
    public string OwnerPassword
    {
      set
      {
        if (this.document.securitySettings.DocumentSecurityLevel == PdfDocumentSecurityLevel.None)
          this.document.securitySettings.DocumentSecurityLevel = PdfDocumentSecurityLevel.Encrypted128Bit;
        this.ownerPassword = value;
      }
    }
    internal string ownerPassword;

    /// <summary>
    /// Gets or sets the user access permission represented as an integer in the P key.
    /// </summary>
    internal PdfUserAccessPermission Permission
    {
      get
      {
        PdfUserAccessPermission permission = (PdfUserAccessPermission)Elements.GetInteger(Keys.P);
        if ((int)permission == 0)
          permission = PdfUserAccessPermission.PermitAll;
        return permission;
      }
      set { Elements.SetInteger(Keys.P, (int)value); }
    }

    /// <summary>
    /// Encrypts the whole document.
    /// </summary>
    public void EncryptDocument()
    {
      foreach (PdfReference iref in this.document.irefTable.AllReferences)
      {
        if (!Object.ReferenceEquals(iref.Value, this))
          EncryptObject(iref.Value);
      }
    }

    /// <summary>
    /// Encrypts an indirect object.
    /// </summary>
    internal void EncryptObject(PdfObject value)
    {
      Debug.Assert(value.Reference != null);

      SetHashKey(value.ObjectID);
#if DEBUG
      if (value.ObjectID.ObjectNumber == 10)
        GetType();
#endif

      PdfDictionary dict;
      PdfArray array;
      PdfStringObject str;
      if ((dict = value as PdfDictionary) != null)
        EncryptDictionary(dict);
      else if ((array = value as PdfArray) != null)
        EncryptArray(array);
      else if ((str = value as PdfStringObject) != null)
      {
        if (str.Length != 0)
        {
          byte[] bytes = str.EncryptionValue;
          PrepareKey();
          EncryptRC4(bytes);
          str.EncryptionValue = bytes;
        }
      }
    }

    /// <summary>
    /// Encrypts a dictionary.
    /// </summary>
    void EncryptDictionary(PdfDictionary dict)
    {
      PdfName[] names = dict.Elements.KeyNames;
      foreach (DictionaryEntry item in dict.Elements)
      {
        PdfString value1;
        PdfDictionary value2;
        PdfArray value3;
        if ((value1 = item.Value as PdfString) != null)
          EncryptString(value1);
        else if ((value2 = item.Value as PdfDictionary) != null)
          EncryptDictionary(value2);
        else if ((value3 = item.Value as PdfArray) != null)
          EncryptArray(value3);
      }
      if (dict.Stream != null)
      {
        byte[] bytes = dict.Stream.Value;
        if (bytes.Length != 0)
        {
          PrepareKey();
          EncryptRC4(bytes);
          dict.Stream.Value = bytes;
        }
      }
    }

    /// <summary>
    /// Encrypts an array.
    /// </summary>
    void EncryptArray(PdfArray array)
    {
      int count = array.Elements.Count;
      for (int idx = 0; idx < count; idx++)
      {
        PdfItem item = array.Elements[idx];
        PdfString value1;
        PdfDictionary value2;
        PdfArray value3;
        if ((value1 = item as PdfString) != null)
          EncryptString(value1);
        else if ((value2 = item as PdfDictionary) != null)
          EncryptDictionary(value2);
        else if ((value3 = item as PdfArray) != null)
          EncryptArray(value3);
      }
    }

    /// <summary>
    /// Encrypts a string.
    /// </summary>
    void EncryptString(PdfString value)
    {
      if (value.Length != 0)
      {
        byte[] bytes = value.EncryptionValue;
        PrepareKey();
        EncryptRC4(bytes);
        value.EncryptionValue = bytes;
      }
    }

    /// <summary>
    /// Encrypts an array.
    /// </summary>
    internal byte[] EncryptBytes(byte[] bytes)
    {
      if (bytes != null && bytes.Length != 0)
      {
        PrepareKey();
        EncryptRC4(bytes);
      }
      return bytes;
    }

    #region Encryption Algorithms

    /// <summary>
    /// Checks the password.
    /// </summary>
    /// <param name="inputPassword">Password or null if no password is provided.</param>
    /// <returns>
    /// 0: inputPassword is neither user nor owner password.
    /// 1: inputPassword is user password.
    /// 2: inputPassword is owner password.
    /// </returns>
    public int ValidatePassword(string inputPassword)
    {
      // We can handle 40 and 128 bit standard encryption
      string filter = Elements.GetName(PdfSecurityHandler.Keys.Filter);
      int v = Elements.GetInteger(Keys.V);
      if (filter != "/Standard" || !(v >= 1 && v <= 3))
        throw new PdfReaderException(PSSR.UnknownEncryption);

      byte[] documentID = PdfEncoders.RawEncoding.GetBytes(Owner.Internals.FirstDocumentID);
      byte[] oValue = PdfEncoders.RawEncoding.GetBytes(Elements.GetString(Keys.O));
      byte[] uValue = PdfEncoders.RawEncoding.GetBytes(Elements.GetString(Keys.U));
      int pValue = Elements.GetInteger(Keys.P);
      int rValue = Elements.GetInteger(Keys.R);

      byte[] password;
      if (inputPassword == null)
        inputPassword = "";

      bool strongEncryption = rValue == 3;
      int keyLength = strongEncryption ? 16 : 32;

      password = PdfEncoders.RawEncoding.GetBytes(inputPassword);
      InitWidhUserPassword(documentID, inputPassword, oValue, pValue, strongEncryption);

      this.document.SecuritySettings.hasOwnerPermissions = false;

      if (!EqualsKey(uValue, keyLength))
      {
        password = PdfEncoders.RawEncoding.GetBytes(inputPassword);

        //Compare owner password
        InitWidhOwnerPassword(documentID, inputPassword, oValue, pValue, strongEncryption);

        if (!EqualsKey(uValue, keyLength))
        {
          //Compare user password
          InitWidhUserPassword(documentID, inputPassword, oValue, pValue, strongEncryption);
          if (!EqualsKey(uValue, keyLength))
            return 0;
          return 1;
        }
        this.document.SecuritySettings.hasOwnerPermissions = true;
        return 2;
      }
      return 1;
    }

    [Conditional("DEBUG")]
    void DumpBytes(string tag, byte[] bytes)
    {
      string dump = tag + ": ";
      for (int idx = 0; idx < bytes.Length; idx++)
        dump += String.Format("{0:X2}", bytes[idx]);
      Debug.WriteLine(dump);
    }

    /// <summary>
    /// Pads a password to a 32 byte array.
    /// </summary>
    byte[] PadPassword(string password)
    {
      byte[] padded = new byte[32];
      if (password == null)
        Array.Copy(passwordPadding, 0, padded, 0, 32);
      else
      {
        int length = password.Length;
        Array.Copy(PdfEncoders.RawEncoding.GetBytes(password), 0, padded, 0, Math.Min(length, 32));
        if (length < 32)
          Array.Copy(passwordPadding, 0, padded, length, 32 - length);
      }
      return padded;
    }
    static byte[] passwordPadding = // 32 bytes password padding defined by Adobe
    {
      0x28, 0xBF, 0x4E, 0x5E, 0x4E, 0x75, 0x8A, 0x41, 0x64, 0x00, 0x4E, 0x56, 0xFF, 0xFA, 0x01, 0x08,
      0x2E, 0x2E, 0x00, 0xB6, 0xD0, 0x68, 0x3E, 0x80, 0x2F, 0x0C, 0xA9, 0xFE, 0x64, 0x53, 0x69, 0x7A,
    };

    /// <summary>
    /// Generates the user key based on the padded user password.
    /// </summary>
    void InitWidhUserPassword(byte[] documentID, string userPassword, byte[] ownerKey, int permissions, bool strongEncryption)
    {
      InitEncryptionKey(documentID, PadPassword(userPassword), ownerKey, permissions, strongEncryption);
      SetupUserKey(documentID);
    }

    /// <summary>
    /// Generates the user key based on the padded owner password.
    /// </summary>
    void InitWidhOwnerPassword(byte[] documentID, string ownerPassword, byte[] ownerKey, int permissions, bool strongEncryption)
    {
      byte[] userPad = ComputeOwnerKey(ownerKey, PadPassword(ownerPassword), strongEncryption);
      InitEncryptionKey(documentID, userPad, ownerKey, permissions, strongEncryption);
      SetupUserKey(documentID);
    }

    /// <summary>
    /// Computes the padded user password from the padded owner password.
    /// </summary>
    byte[] ComputeOwnerKey(byte[] userPad, byte[] ownerPad, bool strongEncryption)
    {
      byte[] ownerKey = new byte[32];

      byte[] digest = this.md5.ComputeHash(ownerPad);
      if (strongEncryption)
      {
        byte[] mkey = new byte[16];
        // Hash the pad 50 times
        for (int idx = 0; idx < 50; idx++)
          digest = this.md5.ComputeHash(digest);
        Array.Copy(userPad, 0, ownerKey, 0, 32);
        // Encrypt the key
        for (int i = 0; i < 20; i++)
        {
          for (int j = 0; j < mkey.Length; ++j)
            mkey[j] = (byte)(digest[j] ^ i);
          PrepareRC4Key(mkey);
          EncryptRC4(ownerKey);
        }
      }
      else
      {
        PrepareRC4Key(digest, 0, 5);
        EncryptRC4(userPad, ownerKey);
      }
      return ownerKey;
    }

    /// <summary>
    /// Computes the encryption key.
    /// </summary>
    void InitEncryptionKey(byte[] documentID, byte[] userPad, byte[] ownerKey, int permissions, bool strongEncryption)
    {
      this.ownerKey = ownerKey;
      this.encryptionKey = new byte[strongEncryption ? 16 : 5];

      this.md5.Initialize();
      this.md5.TransformBlock(userPad, 0, userPad.Length, userPad, 0);
      this.md5.TransformBlock(ownerKey, 0, ownerKey.Length, ownerKey, 0);

      // Split permission into 4 bytes
      byte[] permission = new byte[4];
      permission[0] = (byte)permissions;
      permission[1] = (byte)(permissions >> 8);
      permission[2] = (byte)(permissions >> 16);
      permission[3] = (byte)(permissions >> 24);
      this.md5.TransformBlock(permission, 0, 4, permission, 0);
      this.md5.TransformBlock(documentID, 0, documentID.Length, documentID, 0);
      this.md5.TransformFinalBlock(permission, 0, 0);
      byte[] digest = this.md5.Hash;
      this.md5.Initialize();
      // Create the hash 50 times (only for 128 bit)
      if (this.encryptionKey.Length == 16)
      {
        for (int idx = 0; idx < 50; idx++)
        {
          digest = this.md5.ComputeHash(digest);
          this.md5.Initialize();
        }
      }
      Array.Copy(digest, 0, this.encryptionKey, 0, this.encryptionKey.Length);
    }

    /// <summary>
    /// Computes the user key.
    /// </summary>
    void SetupUserKey(byte[] documentID)
    {
      if (this.encryptionKey.Length == 16)
      {
        this.md5.TransformBlock(passwordPadding, 0, passwordPadding.Length, passwordPadding, 0);
        this.md5.TransformFinalBlock(documentID, 0, documentID.Length);
        byte[] digest = this.md5.Hash;
        this.md5.Initialize();
        Array.Copy(digest, 0, this.userKey, 0, 16);
        for (int idx = 16; idx < 32; idx++)
          this.userKey[idx] = 0;
        //Encrypt the key
        for (int i = 0; i < 20; i++)
        {
          for (int j = 0; j < this.encryptionKey.Length; j++)
            digest[j] = (byte)(this.encryptionKey[j] ^ i);
          PrepareRC4Key(digest, 0, this.encryptionKey.Length);
          EncryptRC4(this.userKey, 0, 16);
        }
      }
      else
      {
        PrepareRC4Key(this.encryptionKey);
        EncryptRC4(passwordPadding, this.userKey);
      }
    }

    /// <summary>
    /// Prepare the encryption key.
    /// </summary>
    void PrepareKey()
    {
      PrepareRC4Key(this.key, 0, this.keySize);
    }

    /// <summary>
    /// Prepare the encryption key.
    /// </summary>
    void PrepareRC4Key(byte[] key)
    {
      PrepareRC4Key(key, 0, key.Length);
    }

    /// <summary>
    /// Prepare the encryption key.
    /// </summary>
    void PrepareRC4Key(byte[] key, int offset, int length)
    {
      int idx1 = 0;
      int idx2 = 0;
      for (int idx = 0; idx < 256; idx++)
        this.state[idx] = (byte)idx;
      byte tmp;
      for (int idx = 0; idx < 256; idx++)
      {
        idx2 = (key[idx1 + offset] + this.state[idx] + idx2) & 255;
        tmp = this.state[idx];
        this.state[idx] = this.state[idx2];
        this.state[idx2] = tmp;
        idx1 = (idx1 + 1) % length;
      }
    }

    /// <summary>
    /// Encrypts the data.
    /// </summary>
    void EncryptRC4(byte[] data)
    {
      EncryptRC4(data, 0, data.Length, data);
    }

    /// <summary>
    /// Encrypts the data.
    /// </summary>
    void EncryptRC4(byte[] data, int offset, int length)
    {
      EncryptRC4(data, offset, length, data);
    }

    /// <summary>
    /// Encrypts the data.
    /// </summary>
    void EncryptRC4(byte[] inputData, byte[] outputData)
    {
      EncryptRC4(inputData, 0, inputData.Length, outputData);
    }

    /// <summary>
    /// Encrypts the data.
    /// </summary>
    void EncryptRC4(byte[] inputData, int offset, int length, byte[] outputData)
    {
      length += offset;
      int x = 0, y = 0;
      byte b;
      for (int idx = offset; idx < length; idx++)
      {
        x = (x + 1) & 255;
        y = (this.state[x] + y) & 255;
        b = this.state[x];
        this.state[x] = this.state[y];
        this.state[y] = b;
        outputData[idx] = (byte)(inputData[idx] ^ state[(this.state[x] + this.state[y]) & 255]);
      }
    }

    /// <summary>
    /// Checks whether the calculated key correct.
    /// </summary>
    bool EqualsKey(byte[] value, int length)
    {
      for (int idx = 0; idx < length; idx++)
      {
        if (this.userKey[idx] != value[idx])
          return false;
      }
      return true;
    }

    /// <summary>
    /// Set the hash key for the specified object.
    /// </summary>
    internal void SetHashKey(PdfObjectID id)
    {
      byte[] objectId = new byte[5];
      this.md5.Initialize();
      // Split the object number and generation
      objectId[0] = (byte)id.ObjectNumber;
      objectId[1] = (byte)(id.ObjectNumber >> 8);
      objectId[2] = (byte)(id.ObjectNumber >> 16);
      objectId[3] = (byte)id.GenerationNumber;
      objectId[4] = (byte)(id.GenerationNumber >> 8);
      this.md5.TransformBlock(this.encryptionKey, 0, this.encryptionKey.Length, this.encryptionKey, 0);
      this.md5.TransformFinalBlock(objectId, 0, objectId.Length);
      this.key = this.md5.Hash;
      this.md5.Initialize();
      this.keySize = this.encryptionKey.Length + 5;
      if (this.keySize > 16)
        this.keySize = 16;
    }

    /// <summary>
    /// Prepares the security handler for encrypting the document.
    /// </summary>
    public void PrepareEncryption()
    {
      Debug.Assert(this.document.securitySettings.DocumentSecurityLevel != PdfDocumentSecurityLevel.None);
      int permissions = (int)this.Permission;
      bool strongEncryption = this.document.securitySettings.DocumentSecurityLevel == PdfDocumentSecurityLevel.Encrypted128Bit;

      PdfInteger vValue;
      PdfInteger length;
      PdfInteger rValue;

      if (strongEncryption)
      {
        vValue = new PdfInteger(2);
        length = new PdfInteger(128);
        rValue = new PdfInteger(3);
      }
      else
      {
        vValue = new PdfInteger(1);
        length = new PdfInteger(40);
        rValue = new PdfInteger(2);
      }

      if (this.userPassword == null || this.userPassword.Length == 0)
        this.userPassword = "";
      // Use user password twice if no owner password provided.
      if (this.ownerPassword == null || this.ownerPassword.Length == 0)
        this.ownerPassword = this.userPassword;

      // Correct permission bits
      permissions |= (int)(strongEncryption ? (uint)0xfffff0c0 : (uint)0xffffffc0);
      permissions &= unchecked((int)0xfffffffc);

      PdfInteger pValue = new PdfInteger(permissions);

      Debug.Assert(this.ownerPassword.Length > 0, "Empty owner password.");
      byte[] userPad = PadPassword(this.userPassword);
      byte[] ownerPad = PadPassword(this.ownerPassword);

      this.md5.Initialize();
      this.ownerKey = ComputeOwnerKey(userPad, ownerPad, strongEncryption);
      byte[] documentID = PdfEncoders.RawEncoding.GetBytes(this.document.Internals.FirstDocumentID);
      InitWidhUserPassword(documentID, this.userPassword, this.ownerKey, permissions, strongEncryption);

      PdfString oValue = new PdfString(PdfEncoders.RawEncoding.GetString(this.ownerKey));
      PdfString uValue = new PdfString(PdfEncoders.RawEncoding.GetString(this.userKey));

      Elements[Keys.Filter] = new PdfName("/Standard");
      Elements[Keys.V] = vValue;
      Elements[Keys.Length] = length;
      Elements[Keys.R] = rValue;
      Elements[Keys.O] = oValue;
      Elements[Keys.U] = uValue;
      Elements[Keys.P] = pValue;
    }

    /// <summary>
    /// The global encryption key.
    /// </summary>
    byte[] encryptionKey;

    /// <summary>
    /// The message digest algorithm MD5.
    /// </summary>
    MD5 md5 = new MD5CryptoServiceProvider();

    /// <summary>
    /// Bytes used for RC4 encryption.
    /// </summary>
    byte[] state = new byte[256];

    /// <summary>
    /// The encryption key for the owner.
    /// </summary>
    byte[] ownerKey = new byte[32];

    /// <summary>
    /// The encryption key for the user.
    /// </summary>
    byte[] userKey = new byte[32];

    /// <summary>
    /// The encryption key for a particular object/generation.
    /// </summary>
    byte[] key;

    /// <summary>
    /// The encryption key length for a particular object/generation.
    /// </summary>
    int keySize;

    #endregion

    internal override void WriteObject(PdfWriter writer)
    {
      // Don't encypt myself
      PdfStandardSecurityHandler securityHandler = writer.SecurityHandler;
      writer.SecurityHandler = null;
      base.WriteObject(writer);
      writer.SecurityHandler = securityHandler;
    }

    #region Keys
    /// <summary>
    /// Predefined keys of this dictionary.
    /// </summary>
    internal sealed new class Keys : PdfSecurityHandler.Keys
    {
      /// <summary>
      /// (Required) A number specifying which revision of the standard security handler
      /// should be used to interpret this dictionary:
      ///  2 if the document is encrypted with a V value less than 2 and does not have any of
      ///   the access permissions set (by means of the P entry, below) that are designated 
      ///   "Revision 3 or greater".
      ///  3 if the document is encrypted with a V value of 2 or 3, or has any "Revision 3 or 
      ///   greater" access permissions set.
      ///  4 if the document is encrypted with a V value of 4
      /// </summary>
      [KeyInfo(KeyType.Integer | KeyType.Required)]
      public const string R = "/R";

      /// <summary>
      /// (Required) A 32-byte string, based on both the owner and user passwords, that is
      /// used in computing the encryption key and in determining whether a valid owner
      /// password was entered.
      /// </summary>
      [KeyInfo(KeyType.String | KeyType.Required)]
      public const string O = "/O";

      /// <summary>
      /// (Required) A 32-byte string, based on the user password, that is used in determining
      /// whether to prompt the user for a password and, if so, whether a valid user or owner 
      /// password was entered.
      /// </summary>
      [KeyInfo(KeyType.String | KeyType.Required)]
      public const string U = "/U";

      /// <summary>
      /// (Required) A set of flags specifying which operations are permitted when the document
      /// is opened with user access.
      /// </summary>
      [KeyInfo(KeyType.Integer | KeyType.Required)]
      public const string P = "/P";

      /// <summary>
      /// (Optional; meaningful only when the value of V is 4; PDF 1.5) Indicates whether
      /// the document-level metadata stream is to be encrypted. Applications should respect this value.
      /// Default value: true.
      /// </summary>
      [KeyInfo(KeyType.Boolean | KeyType.Optional)]
      public const string EncryptMetadata = "/EncryptMetadata";

      /// <summary>
      /// Gets the KeysMeta for these keys.
      /// </summary>
      public static DictionaryMeta Meta
      {
        get
        {
          if (Keys.meta == null)
            Keys.meta = CreateMeta(typeof(Keys));
          return Keys.meta;
        }
      }
      static DictionaryMeta meta;
    }

    /// <summary>
    /// Gets the KeysMeta of this dictionary type.
    /// </summary>
    internal override DictionaryMeta Meta
    {
      get { return Keys.Meta; }
    }
    #endregion
  }
}