Repository: J-F-Liu/lopdf (language: rust)

Issue:
Support RC4 encryption
https://developer.amd.com/wordpress/media/2012/12/AMD_Southern_Islands_Instruction_Set_Architecture.pdf is an example of an encrypted PDF that uses RC4


Discussion / hints:

Fix patch:
diff --git a/Cargo.toml b/Cargo.toml
index 24374cc..9abfbf8 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -26,6 +26,7 @@ image = { version = "^0.24", optional = true }
 itoa = "^1.0"
 linked-hash-map = "^0.5"
 log = "^0.4"
+md5 = "0.7.0"
 nom = { version = "^7.1", optional = true }
 pom = { version = "^3.2", optional = true }
 rayon = { version = "^1.6", optional = true }
diff --git a/examples/extract_text.rs b/examples/extract_text.rs
index c1ceefd..80d7c58 100644
--- a/examples/extract_text.rs
+++ b/examples/extract_text.rs
@@ -63,6 +63,10 @@ pub struct Args {
     /// Optional pretty print output.
     #[clap(short, long)]
     pub pretty: bool,
+
+    /// Optional password for encrypted PDFs
+    #[clap(long, default_value_t = String::from(""))]
+    pub password: String,
 }
 
 impl Args {
@@ -134,11 +138,14 @@ fn get_pdf_text(doc: &Document) -> Result<PdfText, Error> {
     Ok(pdf_text)
 }
 
-fn pdf2text<P: AsRef<Path> + Debug>(path: P, output: P, pretty: bool) -> Result<(), Error> {
+fn pdf2text<P: AsRef<Path> + Debug>(path: P, output: P, pretty: bool, password: &str) -> Result<(), Error> {
     println!("Load {path:?}");
-    let doc = load_pdf(&path)?;
+    let mut doc = load_pdf(&path)?;
     if doc.is_encrypted() {
-        return Err(Error::new(ErrorKind::InvalidInput, "Password missing!"));
+        doc.decrypt(password)
+            .map_err(|_err|
+                Error::new(ErrorKind::InvalidInput, "Failed to decrypt")
+            )?;
     }
     let text = get_pdf_text(&doc)?;
     if !text.errors.is_empty() {
@@ -168,7 +175,7 @@ fn main() -> Result<(), Error> {
     };
     let mut output = PathBuf::from(shellexpand::full(output.to_str().unwrap()).unwrap().to_string());
     output.set_extension("text");
-    pdf2text(&pdf_path, &output, args.pretty)?;
+    pdf2text(&pdf_path, &output, args.pretty, &args.password)?;
     println!(
         "Done after {:.1} seconds.",
         Instant::now().duration_since(start_time).as_secs_f64()
diff --git a/src/document.rs b/src/document.rs
index f597390..e4477f3 100644
--- a/src/document.rs
+++ b/src/document.rs
@@ -1,5 +1,6 @@
 use super::encodings::{self, bytes_to_string, string_to_bytes};
 use super::{Bookmark, Dictionary, Object, ObjectId};
+use crate::encryption;
 use crate::xref::{Xref, XrefType};
 use crate::{Error, Result, Stream};
 use encoding_rs::UTF_16BE;
@@ -253,6 +254,56 @@ impl Document {
         self.get_encrypted().is_ok()
     }
 
+    /// Replaces all encrypted Strings and Streams with their decrypted contents
+    pub fn decrypt<P: AsRef<[u8]>>(&mut self, password: P) -> Result<()> {
+
+        // Find the ID of the encryption dict; we'll want to skip it when decrypting
+        let encryption_obj_id = self.trailer
+            .get(b"Encrypt")
+            .and_then(Object::as_reference)?;
+
+        // Since PDF 1.5, metadata may or may not be encrypted; defaults to true
+        let metadata_is_encrypted = self.get_object(encryption_obj_id)?
+            .as_dict()?
+            .get(b"EncryptMetadata")
+            .and_then(|o| o.as_bool())
+            .unwrap_or(true);
+
+        let key = encryption::get_encryption_key(self, &password, true)?;
+        for (&id, obj) in self.objects.iter_mut() {
+
+            // The encryption dictionary is not encrypted, leave it alone
+            if id == encryption_obj_id {
+                continue;
+            }
+
+            // If a Metadata stream but metadata isn't encrypted, leave it alone
+            if obj.type_name().unwrap_or("") == "Metadata" && !metadata_is_encrypted {
+                continue;
+            }
+
+            let decrypted = match encryption::decrypt_object(&key, id, &*obj) {
+                Ok(content) => content,
+                Err(encryption::DecryptionError::NotDecryptable) => {
+                    continue;
+                }
+                Err(_err) => {
+                    return Err(_err.into());
+                }
+            };
+
+            // Only strings and streams are encrypted
+            match obj {
+                Object::Stream(stream) => stream.set_content(decrypted),
+                Object::String(ref mut content, _) => *content = decrypted,
+                _ => {}
+            }
+        }
+
+        self.trailer.remove(b"Encrypt");
+        Ok(())
+    }
+
     /// Return the PDF document catalog, which is the root of the document's object graph.
     pub fn catalog(&self) -> Result<&Dictionary> {
         self.trailer
diff --git a/src/encryption.rs b/src/encryption.rs
new file mode 100644
index 0000000..e814051
--- /dev/null
+++ b/src/encryption.rs
@@ -0,0 +1,257 @@
+use std::fmt;
+use crate::{Document, Object, ObjectId};
+use crate::rc4::Rc4;
+
+#[derive(Debug)]
+pub enum DecryptionError {
+    MissingEncryptDictionary,
+    MissingRevision,
+    MissingOwnerPassword,
+    MissingPermissions,
+    MissingFileID,
+
+    InvalidKeyLength,
+    InvalidRevision,
+    // Used generically when the object type violates the spec
+    InvalidType,
+
+    NotDecryptable,
+    IncorrectPassword,
+
+    UnsupportedEncryption,
+}
+
+impl std::error::Error for DecryptionError {}
+
+impl fmt::Display for DecryptionError {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        match self {
+            DecryptionError::MissingEncryptDictionary => write!(f, "the /Encrypt dictionary is missing"),
+            DecryptionError::MissingRevision => write!(f, "missing encryption revision"),
+            DecryptionError::MissingOwnerPassword => write!(f, "missing the owner password (/O)"),
+            DecryptionError::MissingPermissions => write!(f, "missing the permissions field (/P)"),
+            DecryptionError::InvalidKeyLength => write!(f, "unsupported key length"),
+            DecryptionError::InvalidRevision => write!(f, "unsupported revision"),
+            DecryptionError::InvalidType => write!(f, "unexpected type; document does not comply with the spec"),
+            DecryptionError::MissingFileID => write!(f, "missing the file /ID elements"),
+            DecryptionError::NotDecryptable => write!(f, "the object is not capable of being decrypted"),
+            DecryptionError::IncorrectPassword => write!(f, "the supplied password is incorrect"),
+            DecryptionError::UnsupportedEncryption=> write!(f, "the document uses an encryption scheme that is not supported"),
+        }
+    }
+}
+
+const PAD_BYTES: [u8; 32] = [
+    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,
+];
+
+const DEFAULT_KEY_LEN: Object = Object::Integer(40);
+const DEFAULT_ALGORITHM: Object = Object::Integer(0);
+
+/// Generates the encryption key for the document and, if `check_password` is
+///  true, verifies that the key is correct.
+pub fn get_encryption_key<P>(doc: &Document, password: P, check_password: bool) -> Result<Vec<u8>, DecryptionError>
+where
+    P: AsRef<[u8]>,
+{
+    let password = password.as_ref();
+
+    let encryption_dict = doc.get_encrypted()
+        .map_err(|_| DecryptionError::MissingEncryptDictionary)?;
+
+    // Very early versions of PDF assume a key length of 40 bits
+    let key_len = encryption_dict.get(b"Length")
+                    .unwrap_or(&DEFAULT_KEY_LEN)
+                    .as_i64()
+                    .map_err(|_| DecryptionError::InvalidType)?
+                    as usize / 8; // Length is in bits, convert to bytes
+
+    // MD5 produces 128bit digests, so key_len must not be greater
+    if key_len > 128/8 {
+        return Err(DecryptionError::InvalidKeyLength);
+    }
+
+    // Make sure we support the encryption algorithm
+    let algorithm = encryption_dict.get(b"V")
+                    .unwrap_or(&DEFAULT_ALGORITHM)
+                    .as_i64()
+                    .map_err(|_| DecryptionError::InvalidType)?;
+    // Currently only support V = 1 or 2
+    if !(1..=2).contains(&algorithm) {
+        return Err(DecryptionError::UnsupportedEncryption);
+    }
+
+    // Revision number dictates hashing strategy
+    let revision = encryption_dict.get(b"R")
+                    .map_err(|_| DecryptionError::MissingRevision)?
+                    .as_i64()
+                    .map_err(|_| DecryptionError::InvalidType)?;
+    if !(2..=3).contains(&revision) {
+        return Err(DecryptionError::UnsupportedEncryption);
+    }
+
+    // Algorithm 3.2
+
+    // 3.2.1 Start building up the key, starting with the user password plaintext,
+    //  padding as needed to 32 bytes
+    let mut key = Vec::with_capacity(128);
+    let password_len = std::cmp::min(password.len(), 32);
+    key.extend_from_slice(&password[0..password_len]);
+    key.extend_from_slice(&PAD_BYTES[0..32-password_len]);
+
+    // 3.2.3 Append hashed owner password
+    let hashed_owner_password = encryption_dict.get(b"O")
+        .map_err(|_| DecryptionError::MissingOwnerPassword)?
+        .as_str()
+        .map_err(|_| DecryptionError::InvalidType)?;
+    key.extend_from_slice(hashed_owner_password);
+
+    // 3.2.4 Append the permissions (4 bytes)
+    let permissions = encryption_dict.get(b"P")
+        // we don't actually care about the permissions but we need the correct
+        //  value to get the correct key
+        .map_err(|_| DecryptionError::MissingPermissions)?
+        .as_i64()
+        .map_err(|_| DecryptionError::InvalidType)? as u32;
+    key.extend_from_slice(&permissions.to_le_bytes());
+
+    // 3.2.5 Append the first element of the file identifier
+    let file_id_0 = doc.trailer.get(b"ID")
+        .map_err(|_| DecryptionError::MissingFileID)?
+        .as_array()
+        .map_err(|_| DecryptionError::InvalidType)?
+        .first().ok_or(DecryptionError::InvalidType)?
+        .as_str()
+        .map_err(|_| DecryptionError::InvalidType)?;
+    key.extend_from_slice(file_id_0);
+
+    // 3.2.6 Revision >=4
+    if revision >= 4 {
+        key.extend_from_slice(&[0xFF_u8, 0xFF, 0xFF, 0xFF]);
+    }
+
+    // 3.2.7+8
+    // Hash the contents of key and take the first key_len bytes
+    let n_hashes = if revision < 3 { 1 } else { 51 };
+    for _ in 0..n_hashes {
+        let digest = md5::compute(&key);
+        key.truncate(key_len); // only keep the first key_len bytes
+        key.copy_from_slice(&digest[..key_len]);
+    }
+
+    // Check that the password is correct
+    if check_password {
+        let check = compute_user_password(&key, revision, file_id_0);
+        if let Ok(Object::String(expected, _)) = encryption_dict.get(b"U") {
+            // Only first 16 bytes are significant, the rest are arbitrary padding
+            if expected[..16] != check[..16] {
+                return Err(DecryptionError::IncorrectPassword);
+            }
+        }
+    }
+
+    Ok(key)
+}
+
+fn compute_user_password<K, ID>(key: K, revision: i64, file_id_0: ID) -> Vec<u8>
+where
+    K: AsRef<[u8]>,
+    ID: AsRef<[u8]>,
+{
+    let key = key.as_ref();
+    let encryptor = Rc4::new(key);
+
+    if revision == 2 {
+        // Algorithm 3.4
+        encryptor.decrypt(PAD_BYTES)
+
+    } else {
+        // Algorithm 3.5
+        // 3.5.2
+        let mut ctx = md5::Context::new();
+        ctx.consume(PAD_BYTES);
+
+        // 3.5.3
+        ctx.consume(file_id_0);
+        let hash = ctx.compute();
+
+        // 3.5.4
+        let mut encrypted_hash = encryptor.encrypt(&hash[..]);
+
+        // 3.5.5
+        let mut temp_key = Vec::new();
+        temp_key.resize(key.len(), 0_u8);
+        for i in 1..=19 {
+            for (in_byte, out_byte) in key.iter().zip(temp_key.iter_mut()) {
+                *out_byte = in_byte ^ (i as u8);
+            }
+            encrypted_hash = Rc4::new(&temp_key).encrypt(encrypted_hash);
+        }
+
+        // 3.5.6
+        encrypted_hash.extend_from_slice(&PAD_BYTES[0..16]);
+
+        encrypted_hash
+    }
+}
+
+/// Decrypts `obj` and returns the content of the string or stream.
+/// If obj is not an decryptable type, returns the NotDecryptable error.
+pub fn decrypt_object<Key>(key: Key, obj_id: ObjectId, obj: &Object) -> Result<Vec<u8>, DecryptionError>
+where
+    Key: AsRef<[u8]>,
+{
+    let key = key.as_ref();
+    let mut builder = Vec::<u8>::with_capacity(key.len() + 5);
+    builder.extend_from_slice(key.as_ref());
+
+    // Extend the key with the lower 3 bytes of the object number
+    builder.extend_from_slice(&obj_id.0.to_le_bytes()[..3]);
+    // and the lower 2 bytes of the generation number
+    builder.extend_from_slice(&obj_id.1.to_le_bytes()[..2]);
+
+    // Now construct the rc4 key
+    let key_len = std::cmp::min(key.len() + 5, 16);
+    let rc4_key = &md5::compute(builder)[..key_len];
+
+    let encrypted = match obj {
+        Object::String(content, _) => content,
+        Object::Stream(stream) => &stream.content,
+        _ => { return Err(DecryptionError::NotDecryptable); },
+    };
+
+    // Decrypt using the rc4 algorithm
+    Ok(Rc4::new(rc4_key).decrypt(encrypted))
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn rc4_works() {
+        let cases = [
+            // Key, Plain, Cipher
+            (String::from("Key"), String::from("Plaintext"), String::from("BBF316E8D940AF0AD3")),
+            (String::from("Wiki"), String::from("pedia"), String::from("1021BF0420")),
+        ];
+
+        for (key, plain, cipher) in cases {
+            // Reencode cipher from a hex string to a Vec<u8>
+            let cipher = cipher.as_bytes();
+            let mut cipher_bytes = Vec::with_capacity(cipher.len()/2);
+            for hex_pair in cipher.chunks_exact(2) {
+                cipher_bytes.push(
+                    u8::from_str_radix(std::str::from_utf8(hex_pair).unwrap(), 16).unwrap()
+                );
+            }
+
+            let decryptor = Rc4::new(key);
+            let decrypted = decryptor.decrypt(&cipher_bytes);
+            assert_eq!(plain.as_bytes(), &decrypted[..]);
+        }
+    }
+}
diff --git a/src/error.rs b/src/error.rs
index e03d391..18eaefc 100644
--- a/src/error.rs
+++ b/src/error.rs
@@ -1,4 +1,5 @@
 use std::fmt;
+use crate::encryption;
 
 #[derive(Debug)]
 pub enum Error {
@@ -44,6 +45,8 @@ pub enum Error {
     /// Error when handling images.
     #[cfg(feature = "embed_image")]
     Image(image::ImageError),
+    /// Error when decrypting the contents of the file
+    Decryption(encryption::DecryptionError),
 }
 
 impl fmt::Display for Error {
@@ -69,6 +72,7 @@ impl fmt::Display for Error {
             Error::NoOutlines => write!(f, "PDF document has no Outlines"),
             #[cfg(feature = "embed_image")]
             Error::Image(e) => e.fmt(f),
+            Error::Decryption(d) => d.fmt(f),
         }
     }
 }
@@ -126,3 +130,9 @@ impl From<image::ImageError> for Error {
         Error::Image(err)
     }
 }
+
+impl From<encryption::DecryptionError> for Error {
+    fn from(_err: encryption::DecryptionError) -> Self {
+        Error::Decryption(_err)
+    }
+}
diff --git a/src/lib.rs b/src/lib.rs
index a34cadd..be022a2 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -26,6 +26,7 @@ pub use crate::toc::Toc;
 pub mod content;
 mod creator;
 mod encodings;
+pub mod encryption;
 mod error;
 pub use error::XrefError;
 pub mod filters;
@@ -39,6 +40,7 @@ mod parser_aux;
 mod processor;
 mod reader;
 pub use reader::Reader;
+mod rc4;
 mod writer;
 pub mod xobject;
 pub use error::{Error, Result};
diff --git a/src/rc4.rs b/src/rc4.rs
new file mode 100644
index 0000000..f3ade8b
--- /dev/null
+++ b/src/rc4.rs
@@ -0,0 +1,69 @@
+// This module exists because the rust-crypto module is really old and not maintained.
+// Fortunately the RC4 algorithm is very simple to implement.
+
+///
+pub struct Rc4 {
+    initial_state: [u8; 256],
+}
+
+impl Rc4 {
+    ///
+    pub fn new<Key: AsRef<[u8]>>(key: Key) -> Self {
+        let key = key.as_ref();
+        assert!(!key.is_empty() && key.len() <= 256);
+
+        let mut initial_state = [0_u8; 256];
+        for (i, v) in initial_state.iter_mut().enumerate() {
+            *v = i as u8;
+        }
+
+        let mut j = 0_u8;
+        for i in 0..256 {
+            j = j.wrapping_add(initial_state[i]).wrapping_add(key[i % key.len()]);
+            initial_state.swap(i, j as usize);
+        }
+
+        Self { initial_state }
+    }
+
+    /// Encrypts/decrypts `input` into `output`.  The shorter of `input` and `output`
+    ///  determine how many bytes are written into `output`.
+    pub fn apply_keystream<'i, 'o, Input, Output>(&self, input: Input, output: Output)
+    where
+        Input: Iterator<Item = &'i u8>,
+        Output: Iterator<Item = &'o mut u8>,
+    {
+        let mut state = self.initial_state;
+        let mut i = 0_u8;
+        let mut j = 0_u8;
+        for (i_byte, o_byte) in input.zip(output) {
+            i = i.wrapping_add(1);
+            j = j.wrapping_add(state[i as usize]);
+            state.swap(i as usize, j as usize);
+            let key_byte = state[(state[i as usize].wrapping_add(state[j as usize])) as usize];
+            *o_byte = i_byte ^ key_byte;
+        }
+    }
+
+    /// Allocates a new Vec<u8> of the same length as `input` and decrypts
+    ///  `input` into it.
+    pub fn decrypt<Input>(&self, input: Input) -> Vec<u8>
+    where
+        Input: AsRef<[u8]>,
+    {
+        let input = input.as_ref();
+        let mut output = vec![0; input.len()];
+        self.apply_keystream(input.iter(), output.iter_mut());
+        output
+    }
+
+    /// Allocates a new Vec<u8> of the same length as `input` and encrypts
+    ///  `input` into it.
+    pub fn encrypt<Input>(&self, input: Input) -> Vec<u8>
+    where
+        Input: AsRef<[u8]>,
+    {
+        // Rc4 is symmetric
+        self.decrypt(input)
+    }
+}


Test patch: