Merge pull request #182 from ReagentX/feat/cs/support-sticker-effects

Feat/cs/support sticker effects

docs/features.md

@@ -37,7 +37,10 @@ This tool targets the current latest public release for macOS and iMessage. It m
 - Attachments
   - Any type of attachment that can be displayed on the web is embedded in the HTML exports
   - Attachments can be copied to the export directory or referenced in-place
-  - Less-compatible HEIC images are converted to PNG for portable exports
+  - Less-compatible images are converted for portable exports:
+    - Attachment `HEIC` files convert to `JPEG`
+    - Sticker `HEIC` files convert to `PNG`
+    - Sticker `HEICS` files convert to `GIF`
   - Attachments are displayed as
     - File paths in TXT exports
     - Embeds in HTML exports (including `<img>`, `<video>`, and `<audio>`)

imessage-database/src/error/attachment.rs

@@ -0,0 +1,28 @@
+/*!
+ Errors that can happen when parsing message data
+*/
+
+use std::{
+    fmt::{Display, Formatter, Result},
+    io::Error,
+};
+
+/// Errors that can happen when working with attachment table data
+#[derive(Debug)]
+pub enum AttachmentError {
+    FileNotFound(String),
+    Unreadable(String, Error),
+}
+
+impl Display for AttachmentError {
+    fn fmt(&self, fmt: &mut Formatter<'_>) -> Result {
+        match self {
+            AttachmentError::FileNotFound(path) => {
+                write!(fmt, "File not found at location: {path}")
+            }
+            AttachmentError::Unreadable(path, why) => {
+                write!(fmt, "Unable to read file at {path}: {why}")
+            }
+        }
+    }
+}

imessage-database/src/error/mod.rs

@@ -2,6 +2,7 @@
  This module contains types of errors that can happen when parsing iMessage data.
 */
 
+pub mod attachment;
 pub mod message;
 pub mod plist;
 pub mod query_context;

imessage-database/src/message_types/mod.rs

@@ -3,9 +3,10 @@
 */
 
 pub mod app;
+pub mod collaboration;
 pub mod edited;
 pub mod expressives;
 pub mod music;
+pub mod sticker;
 pub mod url;
 pub mod variants;
-pub mod collaboration;

imessage-database/src/message_types/sticker.rs

@@ -0,0 +1,168 @@
+use std::fmt::Display;
+
+/// Bytes for `stickerEffect:type="`
+const STICKER_EFFECT_PREFIX: [u8; 20] = [
+    115, 116, 105, 99, 107, 101, 114, 69, 102, 102, 101, 99, 116, 58, 116, 121, 112, 101, 61, 34,
+];
+/// Bytes for `"/>`
+const STICKER_EFFECT_SUFFIX: [u8; 3] = [34, 47, 62];
+
+#[derive(Debug, PartialEq, Eq)]
+pub enum StickerEffect {
+    /// Sticker sent with no effect
+    Normal,
+    /// Internally referred to as `stroke`
+    Outline,
+    Comic,
+    Puffy,
+    /// Internally referred to as `iridescent`
+    Shiny,
+    Other(String),
+}
+
+impl StickerEffect {
+    pub fn from_exif(sticker_effect_type: &str) -> Self {
+        match sticker_effect_type {
+            "stroke" => Self::Outline,
+            "comic" => Self::Comic,
+            "puffy" => Self::Puffy,
+            "iridescent" => Self::Shiny,
+            other => Self::Other(other.to_owned()),
+        }
+    }
+}
+
+impl Display for StickerEffect {
+    fn fmt(&self, fmt: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match self {
+            StickerEffect::Normal => write!(fmt, "Normal"),
+            StickerEffect::Outline => write!(fmt, "Outline"),
+            StickerEffect::Comic => write!(fmt, "Comic"),
+            StickerEffect::Puffy => write!(fmt, "Puffy"),
+            StickerEffect::Shiny => write!(fmt, "Shiny"),
+            StickerEffect::Other(name) => write!(fmt, "{name}"),
+        }
+    }
+}
+
+impl Default for StickerEffect {
+    fn default() -> Self {
+        Self::Normal
+    }
+}
+
+/// Parse the sticker effect type from the EXIF data of a HEIC blob
+pub fn get_sticker_effect(mut heic_data: Vec<u8>) -> StickerEffect {
+    // Find the start index and drain
+    for idx in 0..heic_data.len() {
+        if idx + STICKER_EFFECT_PREFIX.len() < heic_data.len() {
+            let part = &heic_data[idx..idx + STICKER_EFFECT_PREFIX.len()];
+            if part == STICKER_EFFECT_PREFIX {
+                // Remove the start pattern from the blob
+                heic_data.drain(..idx + STICKER_EFFECT_PREFIX.len());
+                break;
+            }
+        } else {
+            return StickerEffect::Normal;
+        }
+    }
+
+    // Find the end index and truncate
+    for idx in 1..heic_data.len() {
+        if idx >= heic_data.len() - 2 {
+            return StickerEffect::Other("Unknown".to_string());
+        }
+        let part = &heic_data[idx..idx + STICKER_EFFECT_SUFFIX.len()];
+
+        if part == STICKER_EFFECT_SUFFIX {
+            // Remove the end pattern from the string
+            heic_data.truncate(idx);
+            break;
+        }
+    }
+    StickerEffect::from_exif(&String::from_utf8_lossy(&heic_data))
+}
+
+#[cfg(test)]
+mod tests {
+    use std::env::current_dir;
+    use std::fs::File;
+    use std::io::Read;
+
+    use crate::message_types::sticker::{get_sticker_effect, StickerEffect};
+
+    #[test]
+    fn test_parse_sticker_normal() {
+        let sticker_path = current_dir()
+            .unwrap()
+            .as_path()
+            .join("test_data/stickers/no_effect.heic");
+        let mut file = File::open(sticker_path).unwrap();
+        let mut bytes = vec![];
+        file.read_to_end(&mut bytes).unwrap();
+
+        let effect = get_sticker_effect(bytes);
+
+        assert_eq!(effect, StickerEffect::Normal);
+    }
+
+    #[test]
+    fn test_parse_sticker_outline() {
+        let sticker_path = current_dir()
+            .unwrap()
+            .as_path()
+            .join("test_data/stickers/outline.heic");
+        let mut file = File::open(sticker_path).unwrap();
+        let mut bytes = vec![];
+        file.read_to_end(&mut bytes).unwrap();
+
+        let effect = get_sticker_effect(bytes);
+
+        assert_eq!(effect, StickerEffect::Outline);
+    }
+
+    #[test]
+    fn test_parse_sticker_comic() {
+        let sticker_path = current_dir()
+            .unwrap()
+            .as_path()
+            .join("test_data/stickers/comic.heic");
+        let mut file = File::open(sticker_path).unwrap();
+        let mut bytes = vec![];
+        file.read_to_end(&mut bytes).unwrap();
+
+        let effect = get_sticker_effect(bytes);
+
+        assert_eq!(effect, StickerEffect::Comic);
+    }
+
+    #[test]
+    fn test_parse_sticker_puffy() {
+        let sticker_path = current_dir()
+            .unwrap()
+            .as_path()
+            .join("test_data/stickers/puffy.heic");
+        let mut file = File::open(sticker_path).unwrap();
+        let mut bytes = vec![];
+        file.read_to_end(&mut bytes).unwrap();
+
+        let effect = get_sticker_effect(bytes);
+
+        assert_eq!(effect, StickerEffect::Puffy);
+    }
+
+    #[test]
+    fn test_parse_sticker_shiny() {
+        let sticker_path = current_dir()
+            .unwrap()
+            .as_path()
+            .join("test_data/stickers/shiny.heic");
+        let mut file = File::open(sticker_path).unwrap();
+        let mut bytes = vec![];
+        file.read_to_end(&mut bytes).unwrap();
+
+        let effect = get_sticker_effect(bytes);
+
+        assert_eq!(effect, StickerEffect::Shiny);
+    }
+}

imessage-database/src/tables/attachment.rs

@@ -4,10 +4,15 @@
 
 use rusqlite::{Connection, Error, Result, Row, Statement};
 use sha1::{Digest, Sha1};
-use std::path::{Path, PathBuf};
+use std::{
+    fs::File,
+    io::Read,
+    path::{Path, PathBuf},
+};
 
 use crate::{
-    error::table::TableError,
+    error::{attachment::AttachmentError, table::TableError},
+    message_types::sticker::{get_sticker_effect, StickerEffect},
     tables::{
         messages::Message,
         table::{Table, ATTACHMENT},
@@ -45,6 +50,7 @@ pub struct Attachment {
     pub mime_type: Option<String>,
     pub transfer_name: Option<String>,
     pub total_bytes: i64,
+    pub is_sticker: bool,
     pub hide_attachment: i32,
     pub copied_path: Option<PathBuf>,
 }
@@ -58,6 +64,7 @@ impl Table for Attachment {
             mime_type: row.get("mime_type").unwrap_or(None),
             transfer_name: row.get("transfer_name").unwrap_or(None),
             total_bytes: row.get("total_bytes").unwrap_or_default(),
+            is_sticker: row.get("is_sticker").unwrap_or(false),
             hide_attachment: row.get("hide_attachment").unwrap_or(0),
             copied_path: None,
         })
@@ -137,6 +144,41 @@ impl Attachment {
         }
     }
 
+    pub fn as_bytes(
+        &self,
+        platform: &Platform,
+        db_path: &Path,
+    ) -> Result<Option<Vec<u8>>, AttachmentError> {
+        if let Some(file_path) = self.resolved_attachment_path(platform, db_path) {
+            let mut file = File::open(&file_path)
+                .map_err(|err| AttachmentError::Unreadable(file_path, err))?;
+            let mut bytes = vec![];
+            file.read_to_end(&mut bytes).unwrap();
+
+            return Ok(Some(bytes));
+        }
+        Ok(None)
+    }
+
+    pub fn get_sticker_effect(
+        &self,
+        platform: &Platform,
+        db_path: &Path,
+    ) -> 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)? {
+            return Ok(Some(get_sticker_effect(data)));
+        }
+
+        // Default if the attachment is a sticker and cannot be parsed/read
+        Ok(Some(StickerEffect::default()))
+    }
+
     /// Get the path to an attachment, if it exists
     pub fn path(&self) -> Option<&Path> {
         match &self.filename {
@@ -331,6 +373,7 @@ mod tests {
             mime_type: Some("image".to_string()),
             transfer_name: Some("c.png".to_string()),
             total_bytes: 100,
+            is_sticker: false,
             hide_attachment: 0,
             copied_path: None,
         }

imessage-exporter/src/app/attachment_manager.rs

@@ -9,7 +9,7 @@ use imessage_database::tables::{attachment::Attachment, messages::Message};
 use uuid::Uuid;
 
 use crate::app::{
-    converter::{heic_to_jpeg, Converter},
+    converter::{convert_heic, Converter, ImageType},
     runtime::Config,
 };
 
@@ -72,7 +72,7 @@ impl AttachmentManager {
             match self {
                 AttachmentManager::Compatible => match &config.converter {
                     Some(converter) => {
-                        Self::copy_convert(from, &mut to, converter);
+                        Self::copy_convert(from, &mut to, converter, attachment.is_sticker);
                     }
                     None => Self::copy_raw(from, &to),
                 },
@@ -116,12 +116,41 @@ impl AttachmentManager {
     }
 
     /// Copy a file, converting if possible
-    fn copy_convert(from: &Path, to: &mut PathBuf, converter: &Converter) {
-        let ext = from.extension().unwrap_or_default();
-        if ext == "heic" || ext == "HEIC" {
+    ///
+    /// - Sticker `HEIC` files convert to `PNG`
+    /// - Sticker `HEICS` files convert to `GIF`
+    /// - Attachment `HEIC` files convert to `JPEG`
+    /// - Other files are just copied with their original formats
+    fn copy_convert(from: &Path, to: &mut PathBuf, converter: &Converter, is_sticker: bool) {
+        let original_extension = from.extension().unwrap_or_default();
+
+        // Handle sticker attachments
+        if is_sticker {
+            // Determine the output type of the sticker
+            let output_type: Option<ImageType> = match original_extension.to_str() {
+                // Normal stickers get converted to png
+                Some("heic") | Some("HEIC") => Some(ImageType::Png),
+                // Animated stickers get converted to gif
+                Some("heics") | Some("HEICS") => Some(ImageType::Gif),
+                _ => None,
+            };
+
+            match output_type {
+                Some(output_type) => {
+                    to.set_extension(output_type.to_str());
+                    if convert_heic(from, to, converter, output_type).is_none() {
+                        eprintln!("Unable to convert {from:?}")
+                    }
+                }
+                None => Self::copy_raw(from, to),
+            }
+        }
+        // Normal attachments always get converted to jpeg
+        else if original_extension == "heic" || original_extension == "HEIC" {
+            let output_type = ImageType::Jpeg;
             // Update extension for conversion
-            to.set_extension("jpg");
-            if heic_to_jpeg(from, to, converter).is_none() {
+            to.set_extension(output_type.to_str());
+            if convert_heic(from, to, converter, output_type).is_none() {
                 eprintln!("Unable to convert {from:?}")
             }
         } else {