Merge pull request #192 from ReagentX/feat/cs/custom-attachment-paths

Feat/cs/custom attachment paths
ReagentX · Nov 19, 2023 · 376a8d4 · 376a8d4
2 parents 1b9508a + 6245004
commit 376a8d4
Show file tree

Hide file tree

Showing 11 changed files with 148 additions and 40 deletions.
diff --git a/Cargo.toml b/Cargo.toml
@@ -1,5 +1,4 @@
 [workspace]
-edition = "2021"
 resolver = "2"
 members = [
   "imessage-database",

diff --git a/imessage-database/src/tables/attachment.rs b/imessage-database/src/tables/attachment.rs
@@ -24,6 +24,7 @@ use crate::{
     },
 };
 
+pub const DEFAULT_ATTACHMENT_ROOT: &str = "~/Library/Messages/Attachments";
 const DIVISOR: f64 = 1024.;
 const UNITS: [&str; 5] = ["B", "KB", "MB", "GB", "TB"];
 
@@ -148,8 +149,11 @@ impl Attachment {
         &self,
         platform: &Platform,
         db_path: &Path,
+        custom_attachment_root: Option<&str>,
     ) -> Result<Option<Vec<u8>>, AttachmentError> {
-        if let Some(file_path) = self.resolved_attachment_path(platform, db_path) {
+        if let Some(file_path) =
+            self.resolved_attachment_path(platform, db_path, custom_attachment_root)
+        {
             let mut file = File::open(&file_path)
                 .map_err(|err| AttachmentError::Unreadable(file_path, err))?;
             let mut bytes = vec![];
@@ -164,14 +168,15 @@ impl Attachment {
         &self,
         platform: &Platform,
         db_path: &Path,
+        custom_attachment_root: Option<&str>,
     ) -> Result<Option<StickerEffect>, AttachmentError> {
         // Handle the non-sticker case
         if !self.is_sticker {
             return Ok(None);
         }
 
         // Try to parse the HEIC data
-        if let Some(data) = self.as_bytes(platform, db_path)? {
+        if let Some(data) = self.as_bytes(platform, db_path, custom_attachment_root)? {
             return Ok(Some(get_sticker_effect(data)));
         }
 
@@ -231,19 +236,31 @@ impl Attachment {
     /// For macOS, `db_path` is unused. For iOS, `db_path` is the path to the root of the backup directory.
     ///
     /// iOS Parsing logic source is from [here](https://github.com/nprezant/iMessageBackup/blob/940d001fb7be557d5d57504eb26b3489e88de26e/imessage_backup_tools.py#L83-L85).
-    pub fn resolved_attachment_path(&self, platform: &Platform, db_path: &Path) -> Option<String> {
-        if let Some(path_str) = &self.filename {
+    ///
+    /// Use the optional `custom_attachment_root` parameter when the attachments are not stored in the same place as the database expects. The expected location is [`DEFAULT_ATTACHMENT_ROOT`](crate::tables::attachment::DEFAULT_ATTACHMENT_ROOT).
+    /// A custom attachment root like `/custom/path` will overwrite a path like `~/Library/Messages/Attachments/3d/...` to `/custom/path/3d...`
+    pub fn resolved_attachment_path(
+        &self,
+        platform: &Platform,
+        db_path: &Path,
+        custom_attachment_root: Option<&str>,
+    ) -> Option<String> {
+        if let Some(mut path_str) = self.filename.clone() {
+            // Apply custom attachment path
+            if let Some(custom_attachment_path) = custom_attachment_root {
+                path_str = path_str.replace(DEFAULT_ATTACHMENT_ROOT, custom_attachment_path);
+            }
             return match platform {
-                Platform::macOS => Some(Attachment::gen_macos_attachment(path_str)),
-                Platform::iOS => Attachment::gen_ios_attachment(path_str, db_path),
+                Platform::macOS => Some(Attachment::gen_macos_attachment(&path_str)),
+                Platform::iOS => Attachment::gen_ios_attachment(&path_str, db_path),
             };
         }
         None
     }
 
     /// Emit diagnostic data for the Attachments table
     ///
-    /// This is defined outside of [crate::tables::table::Diagnostic] because it requires additional data.
+    /// This is defined outside of [`Diagnostic`](crate::tables::table::Diagnostic) because it requires additional data.
     ///
     /// Get the number of attachments that are missing from the filesystem
     /// or are missing one of the following columns:
@@ -359,7 +376,7 @@ impl Attachment {
 #[cfg(test)]
 mod tests {
     use crate::{
-        tables::attachment::{Attachment, MediaType},
+        tables::attachment::{Attachment, MediaType, DEFAULT_ATTACHMENT_ROOT},
         util::platform::Platform,
     };
 
@@ -459,11 +476,24 @@ mod tests {
         let attachment = sample_attachment();
 
         assert_eq!(
-            attachment.resolved_attachment_path(&Platform::macOS, &db_path),
+            attachment.resolved_attachment_path(&Platform::macOS, &db_path, None),
             Some("a/b/c.png".to_string())
         );
     }
 
+    #[test]
+    fn can_get_resolved_path_macos_custom() {
+        let db_path = PathBuf::from("fake_root");
+        let mut attachment = sample_attachment();
+        // Sample path like `~/Library/Messages/Attachments/0a/10/.../image.jpeg`
+        attachment.filename = Some(format!("{DEFAULT_ATTACHMENT_ROOT}/a/b/c.png"));
+
+        assert_eq!(
+            attachment.resolved_attachment_path(&Platform::macOS, &db_path, Some("custom/root")),
+            Some("custom/root/a/b/c.png".to_string())
+        );
+    }
+
     #[test]
     fn can_get_resolved_path_macos_raw() {
         let db_path = PathBuf::from("fake_root");
@@ -472,7 +502,7 @@ mod tests {
 
         assert!(
             attachment
-                .resolved_attachment_path(&Platform::macOS, &db_path)
+                .resolved_attachment_path(&Platform::macOS, &db_path, None)
                 .unwrap()
                 .len()
                 > attachment.filename.unwrap().len()
@@ -486,7 +516,7 @@ mod tests {
         attachment.filename = Some("~/a/b/c~d.png".to_string());
 
         assert!(attachment
-            .resolved_attachment_path(&Platform::macOS, &db_path)
+            .resolved_attachment_path(&Platform::macOS, &db_path, None)
             .unwrap()
             .ends_with("c~d.png"));
     }
@@ -497,7 +527,20 @@ mod tests {
         let attachment = sample_attachment();
 
         assert_eq!(
-            attachment.resolved_attachment_path(&Platform::iOS, &db_path),
+            attachment.resolved_attachment_path(&Platform::iOS, &db_path, None),
+            Some("fake_root/41/41746ffc65924078eae42725c979305626f57cca".to_string())
+        );
+    }
+
+    #[test]
+    fn can_get_resolved_path_ios_custom() {
+        let db_path = PathBuf::from("fake_root");
+        let attachment = sample_attachment();
+
+        // iOS Backups store attachments at the same level as the database file, so if the backup
+        // is intact, the custom root is not relevant
+        assert_eq!(
+            attachment.resolved_attachment_path(&Platform::iOS, &db_path, Some("custom/root")),
             Some("fake_root/41/41746ffc65924078eae42725c979305626f57cca".to_string())
         );
     }
@@ -509,7 +552,7 @@ mod tests {
         attachment.filename = None;
 
         assert_eq!(
-            attachment.resolved_attachment_path(&Platform::macOS, &db_path),
+            attachment.resolved_attachment_path(&Platform::macOS, &db_path, None),
             None
         );
     }
@@ -521,7 +564,7 @@ mod tests {
         attachment.filename = None;
 
         assert_eq!(
-            attachment.resolved_attachment_path(&Platform::iOS, &db_path),
+            attachment.resolved_attachment_path(&Platform::iOS, &db_path, None),
             None
         );
     }

diff --git a/imessage-database/src/tables/messages.rs b/imessage-database/src/tables/messages.rs
@@ -749,7 +749,7 @@ impl Message {
         }
     }
 
-    /// Get the variant of a message, see [crate::message_types::variants] for detail.
+    /// Get the variant of a message, see [`variants`](crate::message_types::variants) for detail.
     pub fn variant(&self) -> Variant {
         // Check if a message was edited first as those have special properties
         if self.is_edited() {

diff --git a/imessage-database/src/tables/table.rs b/imessage-database/src/tables/table.rs
@@ -44,16 +44,25 @@ pub trait Diagnostic {
 
 /// Get a connection to the iMessage SQLite database
 pub fn get_connection(path: &Path) -> Result<Connection, TableError> {
-    if path.exists() {
+    if path.exists() && path.is_file() {
         return match Connection::open_with_flags(path, OpenFlags::SQLITE_OPEN_READ_ONLY) {
             Ok(res) => Ok(res),
             Err(why) => Err(
                 TableError::CannotConnect(
                     format!("Unable to read from chat database: {}\nEnsure full disk access is enabled for your terminal emulator in System Settings > Security and Privacy > Full Disk Access", why)
-                )
-            ),
-        };
+                )),
+            };
+    };
+
+    // Path does not point to a file
+    if path.exists() && !path.is_file() {
+        return Err(TableError::CannotConnect(format!(
+            "Specified path `{}` is not a database!",
+            &path.to_str().unwrap_or("Unknown")
+        )));
     }
+
+    // File is missing
     Err(TableError::CannotConnect(format!(
         "Database not found at {}",
         &path.to_str().unwrap_or("Unknown")

diff --git a/imessage-database/src/util/platform.rs b/imessage-database/src/util/platform.rs
@@ -7,6 +7,7 @@ use std::{fmt::Display, path::Path};
 use crate::tables::table::DEFAULT_PATH_IOS;
 
 /// Represents the platform that created the database this library connects to
+#[derive(PartialEq, Eq)]
 pub enum Platform {
     /// macOS-sourced data
     #[allow(non_camel_case_types)]

diff --git a/imessage-exporter/README.md b/imessage-exporter/README.md
@@ -55,6 +55,11 @@ The [releases page](https://github.com/ReagentX/imessage-exporter/releases) prov
         For iOS, specify a path to the root of an unencrypted backup directory
         If omitted, the default directory is ~/Library/Messages/chat.db
 
+-r, --attachment-root <path/to/attachments>
+        Specify an optional custom path to look for attachments in (macOS only).
+        Only use this if attachments are stored separately from the database's default location.
+        The default location is ~/Library/Messages/Attachments
+
 -a, --platform <macOS, iOS>
         Specify the platform the database was created on
         If omitted, the platform type is determined automatically
@@ -109,6 +114,12 @@ Export as `html` from `/Volumes/external/chat.db` to `/Volumes/external/export`
 % imessage-exporter -f html -c disabled -p /Volumes/external/chat.db -o /Volumes/external/export
 ```
 
+Export as `html` from `/Volumes/external/chat.db` to `/Volumes/external/export` with attachments in `/Volumes/external/Attachments`:
+
+```zsh
+% imessage-exporter -f html -c efficient -p /Volumes/external/chat.db -r /Volumes/external/Attachments -o /Volumes/external/export 
+```
+
 Export messages from `2020-01-01` to `2020-12-31` as `txt` from the default macOS iMessage Database location to `~/export-2020`:
 
 ```zsh

diff --git a/imessage-exporter/src/app/attachment_manager.rs b/imessage-exporter/src/app/attachment_manager.rs
@@ -17,7 +17,7 @@ use crate::app::{
 pub enum AttachmentManager {
     /// Do not copy attachments
     Disabled,
-    /// Copy and convert attachments to more compatible formats using a [crate::app::converter::Converter]
+    /// Copy and convert attachments to more compatible formats using a [`Converter`](crate::app::converter::Converter)
     Compatible,
     /// Copy attachments without converting; preserves quality but may not display correctly in all browsers
     Efficient,
@@ -44,8 +44,11 @@ impl AttachmentManager {
         config: &Config,
     ) -> Option<()> {
         // Resolve the path to the attachment
-        let attachment_path = attachment
-            .resolved_attachment_path(&config.options.platform, &config.options.db_path)?;
+        let attachment_path = attachment.resolved_attachment_path(
+            &config.options.platform,
+            &config.options.db_path,
+            config.options.attachment_root,
+        )?;
 
         if !matches!(self, AttachmentManager::Disabled) {
             let from = Path::new(&attachment_path);