pdfbox/src/main/java/org/apache/pdfbox/pdmodel/encryption/SecurityHandler.java

/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You under the Apache License, Version 2.0
 * (the "License"); you may not use this file except in compliance with
 * the License.  You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package org.apache.pdfbox.pdmodel.encryption;

import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.nio.charset.StandardCharsets;
import java.security.GeneralSecurityException;
import java.security.Key;
import java.security.MessageDigest;
import java.security.SecureRandom;
import java.util.Arrays;
import java.util.Collections;
import java.util.IdentityHashMap;
import java.util.Map;
import java.util.Set;

import javax.crypto.Cipher;
import javax.crypto.CipherInputStream;
import javax.crypto.spec.IvParameterSpec;
import javax.crypto.spec.SecretKeySpec;

import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
import org.apache.pdfbox.cos.COSArray;
import org.apache.pdfbox.cos.COSBase;
import org.apache.pdfbox.cos.COSDictionary;
import org.apache.pdfbox.cos.COSName;
import org.apache.pdfbox.cos.COSStream;
import org.apache.pdfbox.cos.COSString;
import org.apache.pdfbox.io.IOUtils;
import org.apache.pdfbox.pdmodel.PDDocument;

/**
 * A security handler as described in the PDF specifications.
 * A security handler is responsible of documents protection.
 *
 * @author Ben Litchfield
 * @author Benoit Guillon
 * @author Manuel Kasper
 * 
 * @param <T_POLICY> the protection policy.
 */
public abstract class SecurityHandler<T_POLICY extends ProtectionPolicy>
{
    private static final Log LOG = LogFactory.getLog(SecurityHandler.class);

    private static final int DEFAULT_KEY_LENGTH = 40;

    // see 7.6.2, page 58, PDF 32000-1:2008
    private static final byte[] AES_SALT = { (byte) 0x73, (byte) 0x41, (byte) 0x6c, (byte) 0x54 };

    /** The length in bits of the secret key used to encrypt the document. */
    private int keyLength = DEFAULT_KEY_LENGTH;

    /** The encryption key that will used to encrypt / decrypt.*/
    private byte[] encryptionKey;

    /** The RC4 implementation used for cryptographic functions. */
    private final RC4Cipher rc4 = new RC4Cipher();

    /** indicates if the Metadata have to be decrypted of not. */
    private boolean decryptMetadata;

    /** Can be used to allow stateless AES encryption */
    private SecureRandom customSecureRandom;

    // PDFBOX-4453, PDFBOX-4477: Originally this was just a Set. This failed in rare cases
    // when a decrypted string was identical to an encrypted string.
    // Because COSString.equals() checks the contents, decryption was then skipped.
    // This solution keeps all different "equal" objects.
    // IdentityHashMap solves this problem and is also faster than a HashMap
    private final Set<COSBase> objects = Collections.newSetFromMap(new IdentityHashMap<>());

    private boolean useAES;

    /**
     * The typed {@link ProtectionPolicy} to be used for encryption.
     */
    private T_POLICY protectionPolicy = null;
    
    /**
     * The access permission granted to the current user for the document. These
     * permissions are computed during decryption and are in read only mode.
     */
    private AccessPermission currentAccessPermission = null;

    /**
     * The stream filter name.
     */
    private COSName streamFilterName;

    /**
     * The string filter name.
     */
    private COSName stringFilterName;

    /**
     * Constructor.
     */
    protected SecurityHandler()
    {
    }

    /**
     * Constructor used for encryption.
     *
     * @param protectionPolicy The protection policy.
     */
    protected SecurityHandler(T_POLICY protectionPolicy)
    {
        this.protectionPolicy = protectionPolicy;
        keyLength = protectionPolicy.getEncryptionKeyLength();
    }

    /**
     * Set whether to decrypt meta data.
     *
     * @param decryptMetadata true if meta data has to be decrypted.
     */
    protected void setDecryptMetadata(boolean decryptMetadata)
    {
        this.decryptMetadata = decryptMetadata;
    }

    /**
     * Returns true if meta data is to be decrypted.
     *
     * @return True if meta data has to be decrypted.
     */
    public boolean isDecryptMetadata()
    {
        return decryptMetadata;
    }

    /**
     * Set the string filter name.
     * 
     * @param stringFilterName the string filter name.
     */
    protected void setStringFilterName(COSName stringFilterName)
    {
        this.stringFilterName = stringFilterName;
    }

    /**
     * Set the stream filter name.
     * 
     * @param streamFilterName the stream filter name.
     */
    protected void setStreamFilterName(COSName streamFilterName)
    {
        this.streamFilterName = streamFilterName;
    }

    /**
     * Set the custom SecureRandom.
     * 
     * @param customSecureRandom the custom SecureRandom for AES encryption
     */
    public void setCustomSecureRandom(SecureRandom customSecureRandom)
    {
        this.customSecureRandom = customSecureRandom;
    }

    /**
     * Prepare the document for encryption.
     *
     * @param doc The document that will be encrypted.
     *
     * @throws IOException If there is an error with the document.
     */
    public abstract void prepareDocumentForEncryption(PDDocument doc) throws IOException;

    /**
     * Prepares everything to decrypt the document.
     *
     * @param encryption  encryption dictionary, can be retrieved via {@link PDDocument#getEncryption()}
     * @param documentIDArray  document id which is returned via {@link org.apache.pdfbox.cos.COSDocument#getDocumentID()}
     * @param decryptionMaterial Information used to decrypt the document.
     *
     * @throws InvalidPasswordException If the password is incorrect.
     * @throws IOException If there is an error accessing data.
     */
    public abstract void prepareForDecryption(PDEncryption encryption, COSArray documentIDArray,
            DecryptionMaterial decryptionMaterial) throws IOException;

    /**
     * Encrypt or decrypt a set of data.
     *
     * @param objectNumber The data object number.
     * @param genNumber The data generation number.
     * @param data The data to encrypt.
     * @param output The output to write the encrypted data to.
     * @param decrypt true to decrypt the data, false to encrypt it.
     *
     * @throws IOException If there is an error reading the data.
     */
    private void encryptData(long objectNumber, long genNumber, InputStream data,
                            OutputStream output, boolean decrypt) throws IOException
    {
        // Determine whether we're using Algorithm 1 (for RC4 and AES-128), or 1.A (for AES-256)
        if (useAES && encryptionKey.length == 32)
        {
            encryptDataAES256(data, output, decrypt);
        }
        else
        {
            byte[] finalKey = calcFinalKey(objectNumber, genNumber);

            if (useAES)
            {
                encryptDataAESother(finalKey, data, output, decrypt);
            }
            else
            {
                encryptDataRC4(finalKey, data, output);
            }
        }
        output.flush();
    }

    /**
     * Calculate the key to be used for RC4 and AES-128.
     *
     * @param objectNumber The data object number.
     * @param genNumber The data generation number.
     * @return the calculated key.
     */
    private byte[] calcFinalKey(long objectNumber, long genNumber)
    {
        byte[] newKey = new byte[encryptionKey.length + 5];
        System.arraycopy(encryptionKey, 0, newKey, 0, encryptionKey.length);
        // PDF 1.4 reference pg 73
        // step 1
        // we have the reference
        // step 2
        newKey[newKey.length - 5] = (byte) (objectNumber & 0xff);
        newKey[newKey.length - 4] = (byte) (objectNumber >> 8 & 0xff);
        newKey[newKey.length - 3] = (byte) (objectNumber >> 16 & 0xff);
        newKey[newKey.length - 2] = (byte) (genNumber & 0xff);
        newKey[newKey.length - 1] = (byte) (genNumber >> 8 & 0xff);
        // step 3
        MessageDigest md = MessageDigests.getMD5();
        md.update(newKey);
        if (useAES)
        {
            md.update(AES_SALT);
        }
        byte[] digestedKey = md.digest();
        // step 4
        int length = Math.min(newKey.length, 16);
        byte[] finalKey = new byte[length];
        System.arraycopy(digestedKey, 0, finalKey, 0, length);
        return finalKey;
    }

    /**
     * Encrypt or decrypt data with RC4.
     *
     * @param finalKey The final key obtained with via {@link #calcFinalKey(long, long)}.
     * @param input The data to encrypt.
     * @param output The output to write the encrypted data to.
     *
     * @throws IOException If there is an error reading the data.
     */
    protected void encryptDataRC4(byte[] finalKey, InputStream input, OutputStream output)
            throws IOException
    {
        rc4.setKey(finalKey);
        rc4.write(input, output);
    }

    /**
     * Encrypt or decrypt data with RC4.
     *
     * @param finalKey The final key obtained with via {@link #calcFinalKey(long, long)}.
     * @param input The data to encrypt.
     * @param output The output to write the encrypted data to.
     *
     * @throws IOException If there is an error reading the data.
     */
    protected void encryptDataRC4(byte[] finalKey, byte[] input, OutputStream output) throws IOException
    {
        rc4.setKey(finalKey);
        rc4.write(input, output);
    }


    /**
     * Encrypt or decrypt data with AES with key length other than 256 bits.
     *
     * @param finalKey The final key obtained with via {@link #calcFinalKey(long, long)}.
     * @param data The data to encrypt.
     * @param output The output to write the encrypted data to.
     * @param decrypt true to decrypt the data, false to encrypt it.
     *
     * @throws IOException If there is an error reading the data.
     */
    private void encryptDataAESother(byte[] finalKey, InputStream data, OutputStream output, boolean decrypt)
            throws IOException
    {
        byte[] iv = new byte[16];

        if (!prepareAESInitializationVector(decrypt, iv, data, output))
        {
            return;
        }

        try
        {
            Cipher decryptCipher = createCipher(finalKey, iv, decrypt);
            byte[] buffer = new byte[256];
            int n;
            while ((n = data.read(buffer)) != -1)
            {
                byte[] dst = decryptCipher.update(buffer, 0, n);
                if (dst != null)
                {
                    output.write(dst);
                }
            }
            output.write(decryptCipher.doFinal());
        }
        catch (GeneralSecurityException e)
        {
            throw new IOException(e);
        }
    }

    /**
     * Encrypt or decrypt data with AES256.
     *
     * @param data The data to encrypt.
     * @param output The output to write the encrypted data to.
     * @param decrypt true to decrypt the data, false to encrypt it.
     *
     * @throws IOException If there is an error reading the data.
     */
    private void encryptDataAES256(InputStream data, OutputStream output, boolean decrypt) throws IOException
    {
        byte[] iv = new byte[16];

        if (!prepareAESInitializationVector(decrypt, iv, data, output))
        {
            return;
        }

        Cipher cipher;
        try
        {
            cipher = createCipher(this.encryptionKey, iv, decrypt);
        }
        catch (GeneralSecurityException e)
        {
            throw new IOException(e);
        }

        try (CipherInputStream cis = new CipherInputStream(data, cipher))
        {
            IOUtils.copy(cis, output);
        }
        catch (IOException exception)
        {
            // starting with java 8 the JVM wraps an IOException around a GeneralSecurityException
            // it should be safe to swallow a GeneralSecurityException
            if (!(exception.getCause() instanceof GeneralSecurityException))
            {
                throw exception;
            }
            LOG.debug("A GeneralSecurityException occurred when decrypting some stream data", exception);
        }
    }

    private Cipher createCipher(byte[] key, byte[] iv, boolean decrypt) throws GeneralSecurityException
    {
        @SuppressWarnings({"squid:S5542"}) // PKCS#5 padding is requested by PDF specification
        Cipher cipher = Cipher.getInstance("AES/CBC/PKCS5Padding");
        Key keySpec = new SecretKeySpec(key, "AES");
        IvParameterSpec ips = new IvParameterSpec(iv);
        cipher.init(decrypt ? Cipher.DECRYPT_MODE : Cipher.ENCRYPT_MODE, keySpec, ips);
        return cipher;
    }

    private boolean prepareAESInitializationVector(boolean decrypt, byte[] iv, InputStream data, OutputStream output) throws IOException
    {
        if (decrypt)
        {
            // read IV from stream
            int ivSize = (int) IOUtils.populateBuffer(data, iv);
            if (ivSize == 0)
            {
                return false;
            }
            if (ivSize != iv.length)
            {
                throw new IOException(
                        "AES initialization vector not fully read: only "
                                + ivSize + " bytes read instead of " + iv.length);
            }
        }
        else
        {
            // generate random IV and write to stream
            SecureRandom rnd = getSecureRandom();
            rnd.nextBytes(iv);
            output.write(iv);
        }
        return true;
    }

    /**
     * Returns a SecureRandom If customSecureRandom is not defined, instantiate a new SecureRandom
     * 
     * @return SecureRandom
     */
    private SecureRandom getSecureRandom()
    {
        if (customSecureRandom != null)
        {
            return customSecureRandom;
        }
        return new SecureRandom();
    }

    /**
     * This will dispatch to the correct method.
     *
     * @param obj    The object to decrypt.
     * @param objNum The object number.
     * @param genNum The object generation Number.
     *
     * @throws IOException If there is an error getting the stream data.
     */
    public void decrypt(COSBase obj, long objNum, long genNum) throws IOException
    {
        if (!(obj instanceof COSString || obj instanceof COSDictionary || obj instanceof COSArray))
        {
            return;
        }
        // PDFBOX-4477: only cache strings and streams, this improves speed and memory footprint
        if (obj instanceof COSString)
        {
            if (objects.contains(obj))
            {
                return;
            }
            objects.add(obj);
            decryptString((COSString) obj, objNum, genNum);
        }
        else if (obj instanceof COSStream)
        {
            if (objects.contains(obj))
            {
                return;
            }
            objects.add(obj);
            decryptStream((COSStream) obj, objNum, genNum);
        }
        else if (obj instanceof COSDictionary)
        {
            decryptDictionary((COSDictionary) obj, objNum, genNum);
        }
        else if (obj instanceof COSArray)
        {
            decryptArray((COSArray) obj, objNum, genNum);
        }
    }

    /**
     * This will decrypt a stream.
     *
     * @param stream The stream to decrypt.
     * @param objNum The object number.
     * @param genNum The object generation number.
     *
     * @throws IOException If there is an error getting the stream data.
     */
    public void decryptStream(COSStream stream, long objNum, long genNum) throws IOException
    {
        // Stream encrypted with identity filter
        if (COSName.IDENTITY.equals(streamFilterName))
        {
            return;
        }
        
        COSBase type = stream.getCOSName(COSName.TYPE);
        if (!decryptMetadata && COSName.METADATA.equals(type))
        {
            return;
        }
        // "The cross-reference stream shall not be encrypted"
        if (COSName.XREF.equals(type))
        {
            return;
        }
        if (COSName.METADATA.equals(type))
        {
            byte[] buf;
            // PDFBOX-3229 check case where metadata is not encrypted despite /EncryptMetadata missing
            try (InputStream is = stream.createRawInputStream())
            {
                buf = new byte[10];
                long isResult = IOUtils.populateBuffer(is, buf);

                if (Long.compare(isResult, buf.length) != 0)
                {
                    LOG.debug("Tried reading " + buf.length + " bytes but only " + isResult + " bytes read");
                }
            }
            if (Arrays.equals(buf, "<?xpacket ".getBytes(StandardCharsets.ISO_8859_1)))
            {
                LOG.warn("Metadata is not encrypted, but was expected to be");
                LOG.warn("Read PDF specification about EncryptMetadata (default value: true)");
                return;
            }
        }
        decryptDictionary(stream, objNum, genNum);
        byte[] encrypted = IOUtils.toByteArray(stream.createRawInputStream());
        ByteArrayInputStream encryptedStream = new ByteArrayInputStream(encrypted);
        try (OutputStream output = stream.createRawOutputStream())
        {
            encryptData(objNum, genNum, encryptedStream, output, true /* decrypt */);
        }
        catch (IOException ex)
        {
            LOG.error(ex.getClass().getSimpleName() + " thrown when decrypting object " +
                    objNum + " " + genNum + " obj");
            throw ex;
        }
    }

    /**
     * This will encrypt a stream, but not the dictionary as the dictionary is
     * encrypted by visitFromString() in COSWriter and we don't want to encrypt
     * it twice.
     *
     * @param stream The stream to decrypt.
     * @param objNum The object number.
     * @param genNum The object generation number.
     *
     * @throws IOException If there is an error getting the stream data.
     */
    public void encryptStream(COSStream stream, long objNum, int genNum) throws IOException
    {
        // empty streams don't need to be encrypted
        if (!stream.hasData())
        {
            return;
        }
        byte[] rawData = IOUtils.toByteArray(stream.createRawInputStream());
        ByteArrayInputStream encryptedStream = new ByteArrayInputStream(rawData);
        try (OutputStream output = stream.createRawOutputStream())
        {
            encryptData(objNum, genNum, encryptedStream, output, false /* encrypt */);
        }
    }

    /**
     * This will decrypt a dictionary.
     *
     * @param dictionary The dictionary to decrypt.
     * @param objNum The object number.
     * @param genNum The object generation number.
     *
     * @throws IOException If there is an error creating a new string.
     */
    private void decryptDictionary(COSDictionary dictionary, long objNum, long genNum) throws IOException
    {
        if (dictionary.getItem(COSName.CF) != null)
        {
            // PDFBOX-2936: avoid orphan /CF dictionaries found in US govt "I-" files
            return;
        }
        COSName type = dictionary.getCOSName(COSName.TYPE);
        boolean isSignature = COSName.SIG.equals(type) || COSName.DOC_TIME_STAMP.equals(type) ||
                // PDFBOX-4466: /Type is optional, see
                // https://ec.europa.eu/cefdigital/tracker/browse/DSS-1538
                (dictionary.getDictionaryObject(COSName.CONTENTS) instanceof COSString && 
                 dictionary.getDictionaryObject(COSName.BYTERANGE) instanceof COSArray);
        for (Map.Entry<COSName, COSBase> entry : dictionary.entrySet())
        {
            if (isSignature && COSName.CONTENTS.equals(entry.getKey()))
            {
                // do not decrypt the signature contents string
                continue;
            }
            COSBase value = entry.getValue();
            // within a dictionary only the following kind of COS objects have to be decrypted
            if (value instanceof COSString || value instanceof COSArray || value instanceof COSDictionary)
            {
                decrypt(value, objNum, genNum);
            }
        }
    }

    /**
     * This will decrypt a string.
     *
     * @param string the string to decrypt.
     * @param objNum The object number.
     * @param genNum The object generation number.
     *
     * @throws IOException If an error occurs writing the new string.
     */
    private void decryptString(COSString string, long objNum, long genNum) throws IOException
    {
        // String encrypted with identity filter
        if (COSName.IDENTITY.equals(stringFilterName))
        {
            return;
        }
        
        ByteArrayInputStream data = new ByteArrayInputStream(string.getBytes());
        ByteArrayOutputStream outputStream = new ByteArrayOutputStream();
        try
        {
            encryptData(objNum, genNum, data, outputStream, true /* decrypt */);
            string.setValue(outputStream.toByteArray());
        }
        catch (IOException ex)
        {
            LOG.error("Failed to decrypt COSString of length " + string.getBytes().length + 
                    " in object " + objNum + ": " + ex.getMessage(), ex);
        }
    }

    /**
     * This will encrypt a string.
     *
     * @param string the string to encrypt.
     * @param objNum The object number.
     * @param genNum The object generation number.
     *
     * @throws IOException If an error occurs writing the new string.
     */
    public void encryptString(COSString string, long objNum, int genNum) throws IOException
    {
        ByteArrayInputStream data = new ByteArrayInputStream(string.getBytes());
        ByteArrayOutputStream buffer = new ByteArrayOutputStream();
        encryptData(objNum, genNum, data, buffer, false /* encrypt */);
        string.setValue(buffer.toByteArray());
    }

    /**
     * This will decrypt an array.
     *
     * @param array The array to decrypt.
     * @param objNum The object number.
     * @param genNum The object generation number.
     *
     * @throws IOException If there is an error accessing the data.
     */
    private void decryptArray(COSArray array, long objNum, long genNum) throws IOException
    {
        for (int i = 0; i < array.size(); i++)
        {
            decrypt(array.get(i), objNum, genNum);
        }
    }

    /**
     * Getter of the property <tt>keyLength</tt>.
     * @return  Returns the keyLength.
     */
    public int getKeyLength()
    {
        return keyLength;
    }

    /**
     * Setter of the property <tt>keyLength</tt>.
     *
     * @param keyLen  The keyLength to set.
     */
    public void setKeyLength(int keyLen)
    {
        this.keyLength = keyLen;
    }

    /**
     * Sets the access permissions.
     *
     * @param currentAccessPermission The access permissions to be set.
     */
    public void setCurrentAccessPermission(AccessPermission currentAccessPermission)
    {
        this.currentAccessPermission = currentAccessPermission;
    }

    /**
     * Returns the access permissions that were computed during document decryption.
     * The returned object is in read only mode.
     *
     * @return the access permissions or null if the document was not decrypted.
     */
    public AccessPermission getCurrentAccessPermission()
    {
        return currentAccessPermission;
    }

    /**
     * True if AES is used for encryption and decryption.
     *
     * @return true if AEs is used
     */
    public boolean isAES()
    {
        return useAES;
    }

    /**
     * Set to true if AES for encryption and decryption should be used.
     *
     * @param aesValue if true AES will be used
     *
     */
    public void setAES(boolean aesValue)
    {
        useAES = aesValue;
    }

    /**
     * Returns whether a protection policy has been set.
     *
     * @return true if a protection policy has been set.
     */
    public boolean hasProtectionPolicy()
    {
        return protectionPolicy != null;
    }

    /**
     * Returns the set {@link ProtectionPolicy} or null.
     *
     * @return The set {@link ProtectionPolicy}.
     */
    protected T_POLICY getProtectionPolicy()
    {
        return protectionPolicy;
    }

    /**
     * Sets the {@link ProtectionPolicy} to the given value.
     * @param protectionPolicy The {@link ProtectionPolicy}, that shall be set.
     */
    protected void setProtectionPolicy(T_POLICY protectionPolicy)
    {
        this.protectionPolicy = protectionPolicy;
    }

    /**
     * Returns the current encryption key data.
     *
     * @return The current encryption key data.
     */
    public byte[] getEncryptionKey()
    {
        return encryptionKey;
    }

    /**
     * Sets the current encryption key data.
     *
     * @param encryptionKey The encryption key data to set.
     */
    public void setEncryptionKey(byte[] encryptionKey)
    {
        this.encryptionKey = encryptionKey;
    }

    /**
     * Computes the version number of the {@link SecurityHandler} based on the encryption key
     * length. See PDF Spec 1.6 p 93 and
     * <a href="https://www.adobe.com/content/dam/acom/en/devnet/pdf/adobe_supplement_iso32000.pdf">PDF
     * 1.7 Supplement ExtensionLevel: 3</a> and
     * <a href="http://intranet.pdfa.org/wp-content/uploads/2016/08/ISO_DIS_32000-2-DIS4.pdf">PDF
     * Spec 2.0</a>.
     *
     * @return The computed version number.
     */
    protected int computeVersionNumber()
    {
        if (keyLength == 40)
        {
            return 1;
        }
        else if (keyLength == 128 && protectionPolicy.isPreferAES())
        {
            return 4;
        }
        else if (keyLength == 256)
        {
            return 5;
        }
        return 2;
    }
}