Skip to content

Commit

Permalink
Fix #412 and #413
Browse files Browse the repository at this point in the history
  • Loading branch information
@ReagentX
ReagentX committed Dec 25, 2024
1 parent 007e8c4 commit 26a9038
Show file tree
Hide file tree
Showing 2 changed files with 182 additions and 38 deletions.
116 changes: 95 additions & 21 deletions imessage-exporter/src/exporters/html.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -33,7 +33,7 @@ use imessage_database::{
variants::{Announcement, BalloonProvider, CustomBalloon, URLOverride, Variant},
},
tables::{
attachment::{Attachment, MediaType},
attachment::{Attachment, MediaType, StickerSource},
messages::{
models::{AttachmentMeta, BubbleComponent},
Message,
Expand Down Expand Up @@ -596,24 +596,46 @@ impl<'a> Writer<'a> for HTML<'a> {
fn format_sticker(&self, sticker: &'a mut Attachment, message: &Message) -> String {
match self.format_attachment(sticker, message, &AttachmentMeta::default()) {
Ok(mut sticker_embed) => {
// Add sticker effect
let sticker_effect = sticker.get_sticker_effect(
&self.config.options.platform,
&self.config.options.db_path,
self.config.options.attachment_root.as_deref(),
);
if let Ok(Some(sticker_effect)) = sticker_effect {
sticker_embed.push_str(&format!(
"\n<div class=\"sticker_effect\">Sent with {sticker_effect} effect</div>"
))
// Determine the source of the sticker
if let Some(sticker_source) = sticker.get_sticker_source(&self.config.db) {
match sticker_source {
// Emit the prompt used for the Genmoji
StickerSource::Genmoji => {
// Add sticker prompt
if let Some(prompt) = &sticker.emoji_description {
sticker_embed.push_str(&format!(
"\n<div class=\"sticker_effect\">Genmoji prompt: {prompt}</div>"
))
}
}
// Add the app name
StickerSource::Memoji => sticker_embed
.push_str("\n<div class=\"sticker_name\">App: Stickers</div>"),
// Add the sticker effect the user provided
StickerSource::UserGenerated => {
// Add sticker effect
if let Ok(Some(sticker_effect)) = sticker.get_sticker_effect(
&self.config.options.platform,
&self.config.options.db_path,
self.config.options.attachment_root.as_deref(),
) {
sticker_embed.push_str(&format!(
"\n<div class=\"sticker_effect\">Sent with {sticker_effect} effect</div>"
))
}
}
// Add the application name used to generate/send the sticker
StickerSource::App(bundle_id) => {
let app_name = sticker
.get_sticker_source_application_name(&self.config.db)
.unwrap_or(bundle_id);
sticker_embed.push_str(&format!(
"\n<div class=\"sticker_name\">App: {app_name}</div>"
))
}
}
}

// Add sticker prompt
if let Some(prompt) = &sticker.emoji_description {
sticker_embed.push_str(&format!(
"\n<div class=\"sticker_effect\">Genmoji prompt: {prompt}</div>"
))
}
sticker_embed
}
Err(embed) => embed.to_string(),
Expand Down Expand Up @@ -2223,13 +2245,19 @@ mod tests {
// Create exporter
let mut options = Options::fake_options(ExportType::Html);
options.export_path = current_dir().unwrap().parent().unwrap().to_path_buf();
options.db_path = current_dir()
.unwrap()
.parent()
.unwrap()
.join("imessage-database/test_data/db/test.db");

let config = Config::fake_app(options);
let exporter = HTML::new(&config).unwrap();

let message = Config::fake_message();

let mut attachment = Config::fake_attachment();
attachment.rowid = 3;
attachment.is_sticker = true;
let sticker_path = current_dir()
.unwrap()
Expand All @@ -2249,23 +2277,69 @@ mod tests {
.parent()
.unwrap()
.join("orphaned.html");
std::fs::remove_file(orphaned_path).unwrap();
let _ = std::fs::remove_file(orphaned_path);
}

#[test]
fn can_format_html_attachment_sticker_genmoji() {
// Create exporter
let mut options = Options::fake_options(ExportType::Html);
options.export_path = current_dir().unwrap().parent().unwrap().to_path_buf();
options.db_path = current_dir()
.unwrap()
.parent()
.unwrap()
.join("imessage-database/test_data/db/test.db");

let config = Config::fake_app(options);
let exporter = HTML::new(&config).unwrap();

let message = Config::fake_message();

let mut attachment = Config::fake_attachment();
attachment.rowid = 2;
attachment.is_sticker = true;
let sticker_path = current_dir()
.unwrap()
.parent()
.unwrap()
.join("imessage-database/test_data/stickers/outline.heic");
attachment.filename = Some(sticker_path.to_string_lossy().to_string());
attachment.copied_path = Some(PathBuf::from(sticker_path.to_string_lossy().to_string()));
attachment.emoji_description = Some("pink poodle".to_string());

let actual = exporter.format_sticker(&mut attachment, &message);

assert_eq!(actual, "<img src=\"imessage-database/test_data/stickers/outline.heic\" loading=\"lazy\">\n<div class=\"sticker_effect\">Genmoji prompt: pink poodle</div>");

// Remove the file created by the constructor for this test
let orphaned_path = current_dir()
.unwrap()
.parent()
.unwrap()
.join("orphaned.html");
let _ = std::fs::remove_file(orphaned_path);
}

#[test]
fn can_format_html_attachment_sticker_app() {
// Create exporter
let mut options = Options::fake_options(ExportType::Html);
options.export_path = current_dir().unwrap().parent().unwrap().to_path_buf();
options.db_path = current_dir()
.unwrap()
.parent()
.unwrap()
.join("imessage-database/test_data/db/test.db");

let config = Config::fake_app(options);
let exporter = HTML::new(&config).unwrap();

let message = Config::fake_message();

let mut attachment = Config::fake_attachment();
attachment.rowid = 1;
attachment.is_sticker = true;
attachment.emoji_description = Some("Example description".to_string());
let sticker_path = current_dir()
.unwrap()
.parent()
Expand All @@ -2276,15 +2350,15 @@ mod tests {

let actual = exporter.format_sticker(&mut attachment, &message);

assert_eq!(actual, "<img src=\"imessage-database/test_data/stickers/outline.heic\" loading=\"lazy\">\n<div class=\"sticker_effect\">Sent with Outline effect</div>\n<div class=\"sticker_effect\">Genmoji prompt: Example description</div>");
assert_eq!(actual, "<img src=\"imessage-database/test_data/stickers/outline.heic\" loading=\"lazy\">\n<div class=\"sticker_name\">App: Free People</div>");

// Remove the file created by the constructor for this test
let orphaned_path = current_dir()
.unwrap()
.parent()
.unwrap()
.join("orphaned.html");
std::fs::remove_file(orphaned_path).unwrap();
let _ = std::fs::remove_file(orphaned_path);
}

#[test]
Expand Down
104 changes: 87 additions & 17 deletions imessage-exporter/src/exporters/txt.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -34,7 +34,7 @@ use imessage_database::{
variants::{Announcement, BalloonProvider, CustomBalloon, URLOverride, Variant},
},
tables::{
attachment::Attachment,
attachment::{Attachment, StickerSource},
messages::{
models::{AttachmentMeta, BubbleComponent},
Message,
Expand Down Expand Up @@ -415,19 +415,32 @@ impl<'a> Writer<'a> for TXT<'a> {
Ok(path_to_sticker) => {
let mut out_s = format!("Sticker from {who}: {path_to_sticker}");

// Add sticker effect
let sticker_effect = sticker.get_sticker_effect(
&self.config.options.platform,
&self.config.options.db_path,
self.config.options.attachment_root.as_deref(),
);
if let Ok(Some(sticker_effect)) = sticker_effect {
out_s = format!("{sticker_effect} {out_s}");
}

// Add sticker prompt
if let Some(prompt) = &sticker.emoji_description {
out_s = format!("{out_s} (Genmoji prompt: {prompt})");
// Determine the source of the sticker
if let Some(sticker_source) = sticker.get_sticker_source(&self.config.db) {
match sticker_source {
StickerSource::Genmoji => {
if let Some(prompt) = &sticker.emoji_description {
out_s = format!("{out_s} (Genmoji prompt: {prompt})");
}
}
StickerSource::Memoji => out_s.push_str(" (App: Stickers)"),
StickerSource::UserGenerated => {
// Add sticker effect
if let Ok(Some(sticker_effect)) = sticker.get_sticker_effect(
&self.config.options.platform,
&self.config.options.db_path,
self.config.options.attachment_root.as_deref(),
) {
out_s = format!("{sticker_effect} {out_s}");
}
}
StickerSource::App(bundle_id) => {
let app_name = sticker
.get_sticker_source_application_name(&self.config.db)
.unwrap_or(bundle_id);
out_s.push_str(&format!(" (App: {app_name})"));
}
}
}

out_s
Expand Down Expand Up @@ -1661,6 +1674,11 @@ mod tests {
// Create exporter
let mut options = Options::fake_options(ExportType::Txt);
options.export_path = current_dir().unwrap().parent().unwrap().to_path_buf();
options.db_path = current_dir()
.unwrap()
.parent()
.unwrap()
.join("imessage-database/test_data/db/test.db");

let mut config = Config::fake_app(options);
config.participants.insert(0, ME.to_string());
Expand All @@ -1670,6 +1688,7 @@ mod tests {
let message = Config::fake_message();

let mut attachment = Config::fake_attachment();
attachment.rowid = 3;
attachment.is_sticker = true;
let sticker_path = current_dir()
.unwrap()
Expand All @@ -1692,14 +1711,19 @@ mod tests {
.parent()
.unwrap()
.join("orphaned.txt");
std::fs::remove_file(orphaned_path).unwrap();
let _ = std::fs::remove_file(orphaned_path);
}

#[test]
fn can_format_txt_attachment_sticker_genmoji() {
// Create exporter
let mut options = Options::fake_options(ExportType::Txt);
options.export_path = current_dir().unwrap().parent().unwrap().to_path_buf();
options.db_path = current_dir()
.unwrap()
.parent()
.unwrap()
.join("imessage-database/test_data/db/test.db");

let mut config = Config::fake_app(options);
config.participants.insert(0, ME.to_string());
Expand All @@ -1709,6 +1733,7 @@ mod tests {
let message = Config::fake_message();

let mut attachment = Config::fake_attachment();
attachment.rowid = 2;
attachment.is_sticker = true;
attachment.emoji_description = Some("Example description".to_string());
let sticker_path = current_dir()
Expand All @@ -1723,7 +1748,52 @@ mod tests {

assert_eq!(
actual,
"Outline Sticker from Me: imessage-database/test_data/stickers/outline.heic (Genmoji prompt: Example description)"
"Sticker from Me: imessage-database/test_data/stickers/outline.heic (Genmoji prompt: Example description)"
);

// Remove the file created by the constructor for this test
let orphaned_path = current_dir()
.unwrap()
.parent()
.unwrap()
.join("orphaned.txt");
let _ = std::fs::remove_file(orphaned_path);
}

#[test]
fn can_format_txt_attachment_sticker_app() {
// Create exporter
let mut options = Options::fake_options(ExportType::Txt);
options.export_path = current_dir().unwrap().parent().unwrap().to_path_buf();
options.db_path = current_dir()
.unwrap()
.parent()
.unwrap()
.join("imessage-database/test_data/db/test.db");

let mut config = Config::fake_app(options);
config.participants.insert(0, ME.to_string());

let exporter = TXT::new(&config).unwrap();

let message = Config::fake_message();

let mut attachment = Config::fake_attachment();
attachment.rowid = 1;
attachment.is_sticker = true;
let sticker_path = current_dir()
.unwrap()
.parent()
.unwrap()
.join("imessage-database/test_data/stickers/outline.heic");
attachment.filename = Some(sticker_path.to_string_lossy().to_string());
attachment.copied_path = Some(PathBuf::from(sticker_path.to_string_lossy().to_string()));

let actual = exporter.format_sticker(&mut attachment, &message);

assert_eq!(
actual,
"Sticker from Me: imessage-database/test_data/stickers/outline.heic (App: Free People)"
);

// Remove the file created by the constructor for this test
Expand All @@ -1732,7 +1802,7 @@ mod tests {
.parent()
.unwrap()
.join("orphaned.txt");
std::fs::remove_file(orphaned_path).unwrap();
let _ = std::fs::remove_file(orphaned_path);
}

#[test]
Expand Down

0 comments on commit 26a9038

Please sign in to comment.