1use plist::Value;
6use rusqlite::{Connection, Error, Result, Row, Statement, blob::Blob};
7use sha1::{Digest, Sha1};
8
9use std::{
10 fs::File,
11 io::Read,
12 path::{Path, PathBuf},
13};
14
15use crate::{
16 error::{attachment::AttachmentError, table::TableError},
17 message_types::sticker::{StickerEffect, StickerSource, get_sticker_effect},
18 tables::{
19 messages::Message,
20 table::{ATTACHMENT, ATTRIBUTION_INFO, GetBlob, STICKER_USER_INFO, Table},
21 },
22 util::{
23 dates::TIMESTAMP_FACTOR,
24 dirs::home,
25 output::{done_processing, processing},
26 platform::Platform,
27 plist::plist_as_dictionary,
28 query_context::QueryContext,
29 size::format_file_size,
30 },
31};
32
33pub const DEFAULT_ATTACHMENT_ROOT: &str = "~/Library/Messages/Attachments";
35const COLS: &str = "a.rowid, a.filename, a.uti, a.mime_type, a.transfer_name, a.total_bytes, a.is_sticker, a.hide_attachment, a.emoji_image_short_description";
36
37#[derive(Debug, PartialEq, Eq)]
41pub enum MediaType<'a> {
42 Image(&'a str),
44 Video(&'a str),
46 Audio(&'a str),
48 Text(&'a str),
50 Application(&'a str),
52 Other(&'a str),
54 Unknown,
56}
57
58impl MediaType<'_> {
59 #[must_use]
69 pub fn as_mime_type(&self) -> String {
70 match self {
71 MediaType::Image(subtype) => format!("image/{subtype}"),
72 MediaType::Video(subtype) => format!("video/{subtype}"),
73 MediaType::Audio(subtype) => format!("audio/{subtype}"),
74 MediaType::Text(subtype) => format!("text/{subtype}"),
75 MediaType::Application(subtype) => format!("application/{subtype}"),
76 MediaType::Other(mime) => (*mime).to_string(),
77 MediaType::Unknown => String::new(),
78 }
79 }
80}
81
82#[derive(Debug)]
84pub struct Attachment {
85 pub rowid: i32,
87 pub filename: Option<String>,
89 pub uti: Option<String>,
91 pub mime_type: Option<String>,
93 pub transfer_name: Option<String>,
95 pub total_bytes: i64,
97 pub is_sticker: bool,
99 pub hide_attachment: i32,
101 pub emoji_description: Option<String>,
103 pub copied_path: Option<PathBuf>,
105}
106
107impl Table for Attachment {
108 fn from_row(row: &Row) -> Result<Attachment> {
109 Ok(Attachment {
110 rowid: row.get("rowid")?,
111 filename: row.get("filename").unwrap_or(None),
112 uti: row.get("uti").unwrap_or(None),
113 mime_type: row.get("mime_type").unwrap_or(None),
114 transfer_name: row.get("transfer_name").unwrap_or(None),
115 total_bytes: row.get("total_bytes").unwrap_or_default(),
116 is_sticker: row.get("is_sticker").unwrap_or(false),
117 hide_attachment: row.get("hide_attachment").unwrap_or(0),
118 emoji_description: row.get("emoji_image_short_description").unwrap_or(None),
119 copied_path: None,
120 })
121 }
122
123 fn get(db: &Connection) -> Result<Statement, TableError> {
124 Ok(db.prepare(&format!("SELECT * from {ATTACHMENT}"))?)
125 }
126
127 fn extract(attachment: Result<Result<Self, Error>, Error>) -> Result<Self, TableError> {
128 match attachment {
129 Ok(Ok(attachment)) => Ok(attachment),
130 Err(why) | Ok(Err(why)) => Err(TableError::QueryError(why)),
131 }
132 }
133}
134
135impl GetBlob for Attachment {
136 fn get_blob<'a>(&self, db: &'a Connection, column: &str) -> Option<Blob<'a>> {
138 db.blob_open(
139 rusqlite::MAIN_DB,
140 ATTACHMENT,
141 column,
142 i64::from(self.rowid),
143 true,
144 )
145 .ok()
146 }
147}
148
149impl Attachment {
150 pub fn from_message(db: &Connection, msg: &Message) -> Result<Vec<Attachment>, TableError> {
154 let mut out_l = vec![];
155 if msg.has_attachments() {
156 let mut statement = db
157 .prepare(&format!(
158 "
159 SELECT {COLS}
160 FROM message_attachment_join j
161 LEFT JOIN {ATTACHMENT} a ON j.attachment_id = a.ROWID
162 WHERE j.message_id = {}
163 ",
164 msg.rowid
165 ))
166 .or_else(|_| {
167 db.prepare(&format!(
168 "
169 SELECT *
170 FROM message_attachment_join j
171 LEFT JOIN {ATTACHMENT} a ON j.attachment_id = a.ROWID
172 WHERE j.message_id = {}
173 ",
174 msg.rowid
175 ))
176 })?;
177
178 let iter = statement.query_map([], |row| Ok(Attachment::from_row(row)))?;
179
180 for attachment in iter {
181 let m = Attachment::extract(attachment)?;
182 out_l.push(m);
183 }
184 }
185 Ok(out_l)
186 }
187
188 #[must_use]
190 pub fn mime_type(&'_ self) -> MediaType<'_> {
191 match &self.mime_type {
192 Some(mime) => {
193 let mut mime_parts = mime.split('/');
194 if let (Some(category), Some(subtype)) = (mime_parts.next(), mime_parts.next()) {
195 match category {
196 "image" => MediaType::Image(subtype),
197 "video" => MediaType::Video(subtype),
198 "audio" => MediaType::Audio(subtype),
199 "text" => MediaType::Text(subtype),
200 "application" => MediaType::Application(subtype),
201 _ => MediaType::Other(mime),
202 }
203 } else {
204 MediaType::Other(mime)
205 }
206 }
207 None => {
208 if let Some(uti) = &self.uti {
210 match uti.as_str() {
211 "com.apple.coreaudio-format" => MediaType::Audio("x-caf; codecs=opus"),
214 _ => MediaType::Unknown,
215 }
216 } else {
217 MediaType::Unknown
218 }
219 }
220 }
221 }
222
223 pub fn as_bytes(
228 &self,
229 platform: &Platform,
230 db_path: &Path,
231 custom_attachment_root: Option<&str>,
232 ) -> Result<Option<Vec<u8>>, AttachmentError> {
233 if let Some(file_path) =
234 self.resolved_attachment_path(platform, db_path, custom_attachment_root)
235 {
236 let mut file = File::open(&file_path)
237 .map_err(|err| AttachmentError::Unreadable(file_path.clone(), err))?;
238 let mut bytes = vec![];
239 file.read_to_end(&mut bytes)
240 .map_err(|err| AttachmentError::Unreadable(file_path.clone(), err))?;
241
242 return Ok(Some(bytes));
243 }
244 Ok(None)
245 }
246
247 pub fn get_sticker_effect(
252 &self,
253 platform: &Platform,
254 db_path: &Path,
255 custom_attachment_root: Option<&str>,
256 ) -> Result<Option<StickerEffect>, AttachmentError> {
257 if !self.is_sticker {
259 return Ok(None);
260 }
261
262 if let Some(data) = self.as_bytes(platform, db_path, custom_attachment_root)? {
264 return Ok(Some(get_sticker_effect(data)));
265 }
266
267 Ok(Some(StickerEffect::default()))
269 }
270
271 #[must_use]
273 pub fn path(&self) -> Option<&Path> {
274 match &self.filename {
275 Some(name) => Some(Path::new(name)),
276 None => None,
277 }
278 }
279
280 #[must_use]
282 pub fn extension(&self) -> Option<&str> {
283 match self.path() {
284 Some(path) => match path.extension() {
285 Some(ext) => ext.to_str(),
286 None => None,
287 },
288 None => None,
289 }
290 }
291
292 #[must_use]
296 pub fn filename(&self) -> Option<&str> {
297 self.transfer_name.as_deref().or(self.filename.as_deref())
298 }
299
300 #[must_use]
302 pub fn file_size(&self) -> String {
303 format_file_size(u64::try_from(self.total_bytes).unwrap_or(0))
304 }
305
306 pub fn get_total_attachment_bytes(
308 db: &Connection,
309 context: &QueryContext,
310 ) -> Result<u64, TableError> {
311 let mut bytes_query = if context.start.is_some() || context.end.is_some() {
312 let mut statement = format!("SELECT IFNULL(SUM(total_bytes), 0) FROM {ATTACHMENT} a");
313
314 statement.push_str(" WHERE ");
315 if let Some(start) = context.start {
316 statement.push_str(&format!(
317 " a.created_date >= {}",
318 start / TIMESTAMP_FACTOR
319 ));
320 }
321 if let Some(end) = context.end {
322 if context.start.is_some() {
323 statement.push_str(" AND ");
324 }
325 statement.push_str(&format!(" a.created_date <= {}", end / TIMESTAMP_FACTOR));
326 }
327
328 db.prepare(&statement)?
329 } else {
330 db.prepare(&format!(
331 "SELECT IFNULL(SUM(total_bytes), 0) FROM {ATTACHMENT}"
332 ))?
333 };
334 Ok(bytes_query
335 .query_row([], |r| -> Result<i64> { r.get(0) })
336 .map(|res: i64| u64::try_from(res).unwrap_or(0))?)
337 }
338
339 #[must_use]
351 pub fn resolved_attachment_path(
352 &self,
353 platform: &Platform,
354 db_path: &Path,
355 custom_attachment_root: Option<&str>,
356 ) -> Option<String> {
357 if let Some(mut path_str) = self.filename.clone() {
358 if let Some(custom_attachment_path) = custom_attachment_root {
360 path_str = path_str.replace(DEFAULT_ATTACHMENT_ROOT, custom_attachment_path);
361 }
362 return match platform {
363 Platform::macOS => Some(Attachment::gen_macos_attachment(&path_str)),
364 Platform::iOS => Attachment::gen_ios_attachment(&path_str, db_path),
365 };
366 }
367 None
368 }
369
370 pub fn run_diagnostic(
392 db: &Connection,
393 db_path: &Path,
394 platform: &Platform,
395 ) -> Result<(), TableError> {
396 processing();
397 let mut total_attachments = 0;
398 let mut null_attachments = 0;
399 let mut size_on_disk: u64 = 0;
400 let mut statement_paths = db.prepare(&format!("SELECT filename FROM {ATTACHMENT}"))?;
401 let paths = statement_paths.query_map([], |r| Ok(r.get(0)))?;
402
403 let missing_files = paths
404 .filter_map(Result::ok)
405 .filter(|path: &Result<String, Error>| {
406 total_attachments += 1;
408 if let Ok(filepath) = path {
409 match platform {
410 Platform::macOS => {
411 let path = Attachment::gen_macos_attachment(filepath);
412 let file = Path::new(&path);
413 if let Ok(metadata) = file.metadata() {
414 size_on_disk += metadata.len();
415 }
416 !file.exists()
417 }
418 Platform::iOS => {
419 if let Some(parsed_path) =
420 Attachment::gen_ios_attachment(filepath, db_path)
421 {
422 let file = Path::new(&parsed_path);
423 if let Ok(metadata) = file.metadata() {
424 size_on_disk += metadata.len();
425 }
426 return !file.exists();
427 }
428 true
430 }
431 }
432 } else {
433 null_attachments += 1;
435 true
436 }
437 })
438 .count();
439
440 let total_bytes =
441 Attachment::get_total_attachment_bytes(db, &QueryContext::default()).unwrap_or(0);
442
443 done_processing();
444
445 if total_attachments > 0 {
446 println!("\rAttachment diagnostic data:");
447 println!(" Total attachments: {total_attachments}");
448 println!(
449 " Data referenced in table: {}",
450 format_file_size(total_bytes)
451 );
452 println!(
453 " Data present on disk: {}",
454 format_file_size(size_on_disk)
455 );
456 if missing_files > 0 && total_attachments > 0 {
457 println!(
458 " Missing files: {missing_files:?} ({:.0}%)",
459 (missing_files as f64 / f64::from(total_attachments)) * 100f64
460 );
461 println!(" No path provided: {null_attachments}");
462 println!(
463 " No file located: {}",
464 missing_files.saturating_sub(null_attachments)
465 );
466 }
467 }
468 Ok(())
469 }
470
471 fn gen_macos_attachment(path: &str) -> String {
473 if path.starts_with('~') {
474 return path.replacen('~', &home(), 1);
475 }
476 path.to_string()
477 }
478
479 fn gen_ios_attachment(file_path: &str, db_path: &Path) -> Option<String> {
481 let input = file_path.get(2..)?;
482 let filename = format!(
483 "{:x}",
484 Sha1::digest(format!("MediaDomain-{input}").as_bytes())
485 );
486 let directory = filename.get(0..2)?;
487
488 Some(format!("{}/{directory}/{filename}", db_path.display()))
489 }
490
491 fn sticker_info(&self, db: &Connection) -> Option<Value> {
498 Value::from_reader(self.get_blob(db, STICKER_USER_INFO)?).ok()
499 }
500
501 fn attribution_info(&self, db: &Connection) -> Option<Value> {
508 Value::from_reader(self.get_blob(db, ATTRIBUTION_INFO)?).ok()
509 }
510
511 pub fn get_sticker_source(&self, db: &Connection) -> Option<StickerSource> {
516 if let Some(sticker_info) = self.sticker_info(db) {
517 let plist = plist_as_dictionary(&sticker_info).ok()?;
518 let bundle_id = plist.get("pid")?.as_string()?;
519 return StickerSource::from_bundle_id(bundle_id);
520 }
521 None
522 }
523
524 pub fn get_sticker_source_application_name(&self, db: &Connection) -> Option<String> {
529 if let Some(attribution_info) = self.attribution_info(db) {
530 let plist = plist_as_dictionary(&attribution_info).ok()?;
531 return Some(plist.get("name")?.as_string()?.to_owned());
532 }
533 None
534 }
535}
536
537#[cfg(test)]
538mod tests {
539 use crate::{
540 tables::{
541 attachment::{Attachment, DEFAULT_ATTACHMENT_ROOT, MediaType},
542 table::get_connection,
543 },
544 util::{platform::Platform, query_context::QueryContext},
545 };
546
547 use std::{
548 collections::BTreeSet,
549 env::current_dir,
550 path::{Path, PathBuf},
551 };
552
553 fn sample_attachment() -> Attachment {
554 Attachment {
555 rowid: 1,
556 filename: Some("a/b/c.png".to_string()),
557 uti: Some("public.png".to_string()),
558 mime_type: Some("image/png".to_string()),
559 transfer_name: Some("c.png".to_string()),
560 total_bytes: 100,
561 is_sticker: false,
562 hide_attachment: 0,
563 emoji_description: None,
564 copied_path: None,
565 }
566 }
567
568 #[test]
569 fn can_get_path() {
570 let attachment = sample_attachment();
571 assert_eq!(attachment.path(), Some(Path::new("a/b/c.png")));
572 }
573
574 #[test]
575 fn cant_get_path_missing() {
576 let mut attachment = sample_attachment();
577 attachment.filename = None;
578 assert_eq!(attachment.path(), None);
579 }
580
581 #[test]
582 fn can_get_extension() {
583 let attachment = sample_attachment();
584 assert_eq!(attachment.extension(), Some("png"));
585 }
586
587 #[test]
588 fn cant_get_extension_missing() {
589 let mut attachment = sample_attachment();
590 attachment.filename = None;
591 assert_eq!(attachment.extension(), None);
592 }
593
594 #[test]
595 fn can_get_mime_type_png() {
596 let attachment = sample_attachment();
597 assert_eq!(attachment.mime_type(), MediaType::Image("png"));
598 }
599
600 #[test]
601 fn can_get_mime_type_heic() {
602 let mut attachment = sample_attachment();
603 attachment.mime_type = Some("image/heic".to_string());
604 assert_eq!(attachment.mime_type(), MediaType::Image("heic"));
605 }
606
607 #[test]
608 fn can_get_mime_type_fake() {
609 let mut attachment = sample_attachment();
610 attachment.mime_type = Some("fake/bloop".to_string());
611 assert_eq!(attachment.mime_type(), MediaType::Other("fake/bloop"));
612 }
613
614 #[test]
615 fn can_get_mime_type_missing() {
616 let mut attachment = sample_attachment();
617 attachment.mime_type = None;
618 assert_eq!(attachment.mime_type(), MediaType::Unknown);
619 }
620
621 #[test]
622 fn can_get_filename() {
623 let attachment = sample_attachment();
624 assert_eq!(attachment.filename(), Some("c.png"));
625 }
626
627 #[test]
628 fn can_get_filename_no_transfer_name() {
629 let mut attachment = sample_attachment();
630 attachment.transfer_name = None;
631 assert_eq!(attachment.filename(), Some("a/b/c.png"));
632 }
633
634 #[test]
635 fn can_get_filename_no_filename() {
636 let mut attachment = sample_attachment();
637 attachment.filename = None;
638 assert_eq!(attachment.filename(), Some("c.png"));
639 }
640
641 #[test]
642 fn can_get_filename_no_meta() {
643 let mut attachment = sample_attachment();
644 attachment.transfer_name = None;
645 attachment.filename = None;
646 assert_eq!(attachment.filename(), None);
647 }
648
649 #[test]
650 fn can_get_resolved_path_macos() {
651 let db_path = PathBuf::from("fake_root");
652 let attachment = sample_attachment();
653
654 assert_eq!(
655 attachment.resolved_attachment_path(&Platform::macOS, &db_path, None),
656 Some("a/b/c.png".to_string())
657 );
658 }
659
660 #[test]
661 fn can_get_resolved_path_macos_custom() {
662 let db_path = PathBuf::from("fake_root");
663 let mut attachment = sample_attachment();
664 attachment.filename = Some(format!("{DEFAULT_ATTACHMENT_ROOT}/a/b/c.png"));
666
667 assert_eq!(
668 attachment.resolved_attachment_path(&Platform::macOS, &db_path, Some("custom/root")),
669 Some("custom/root/a/b/c.png".to_string())
670 );
671 }
672
673 #[test]
674 fn can_get_resolved_path_macos_raw() {
675 let db_path = PathBuf::from("fake_root");
676 let mut attachment = sample_attachment();
677 attachment.filename = Some("~/a/b/c.png".to_string());
678
679 assert!(
680 attachment
681 .resolved_attachment_path(&Platform::macOS, &db_path, None)
682 .unwrap()
683 .len()
684 > attachment.filename.unwrap().len()
685 );
686 }
687
688 #[test]
689 fn can_get_resolved_path_macos_raw_tilde() {
690 let db_path = PathBuf::from("fake_root");
691 let mut attachment = sample_attachment();
692 attachment.filename = Some("~/a/b/c~d.png".to_string());
693
694 assert!(
695 attachment
696 .resolved_attachment_path(&Platform::macOS, &db_path, None)
697 .unwrap()
698 .ends_with("c~d.png")
699 );
700 }
701
702 #[test]
703 fn can_get_resolved_path_ios() {
704 let db_path = PathBuf::from("fake_root");
705 let attachment = sample_attachment();
706
707 assert_eq!(
708 attachment.resolved_attachment_path(&Platform::iOS, &db_path, None),
709 Some("fake_root/41/41746ffc65924078eae42725c979305626f57cca".to_string())
710 );
711 }
712
713 #[test]
714 fn can_get_resolved_path_ios_custom() {
715 let db_path = PathBuf::from("fake_root");
716 let attachment = sample_attachment();
717
718 assert_eq!(
721 attachment.resolved_attachment_path(&Platform::iOS, &db_path, Some("custom/root")),
722 Some("fake_root/41/41746ffc65924078eae42725c979305626f57cca".to_string())
723 );
724 }
725
726 #[test]
727 fn cant_get_missing_resolved_path_macos() {
728 let db_path = PathBuf::from("fake_root");
729 let mut attachment = sample_attachment();
730 attachment.filename = None;
731
732 assert_eq!(
733 attachment.resolved_attachment_path(&Platform::macOS, &db_path, None),
734 None
735 );
736 }
737
738 #[test]
739 fn cant_get_missing_resolved_path_ios() {
740 let db_path = PathBuf::from("fake_root");
741 let mut attachment = sample_attachment();
742 attachment.filename = None;
743
744 assert_eq!(
745 attachment.resolved_attachment_path(&Platform::iOS, &db_path, None),
746 None
747 );
748 }
749
750 #[test]
751 fn can_get_attachment_bytes_no_filter() {
752 let db_path = current_dir()
753 .unwrap()
754 .parent()
755 .unwrap()
756 .join("imessage-database/test_data/db/test.db");
757 let connection = get_connection(&db_path).unwrap();
758
759 let context = QueryContext::default();
760
761 assert!(Attachment::get_total_attachment_bytes(&connection, &context).is_ok());
762 }
763
764 #[test]
765 fn can_get_attachment_bytes_start_filter() {
766 let db_path = current_dir()
767 .unwrap()
768 .parent()
769 .unwrap()
770 .join("imessage-database/test_data/db/test.db");
771 let connection = get_connection(&db_path).unwrap();
772
773 let mut context = QueryContext::default();
774 context.set_start("2020-01-01").unwrap();
775
776 assert!(Attachment::get_total_attachment_bytes(&connection, &context).is_ok());
777 }
778
779 #[test]
780 fn can_get_attachment_bytes_end_filter() {
781 let db_path = current_dir()
782 .unwrap()
783 .parent()
784 .unwrap()
785 .join("imessage-database/test_data/db/test.db");
786 let connection = get_connection(&db_path).unwrap();
787
788 let mut context = QueryContext::default();
789 context.set_end("2020-01-01").unwrap();
790
791 assert!(Attachment::get_total_attachment_bytes(&connection, &context).is_ok());
792 }
793
794 #[test]
795 fn can_get_attachment_bytes_start_end_filter() {
796 let db_path = current_dir()
797 .unwrap()
798 .parent()
799 .unwrap()
800 .join("imessage-database/test_data/db/test.db");
801 let connection = get_connection(&db_path).unwrap();
802
803 let mut context = QueryContext::default();
804 context.set_start("2020-01-01").unwrap();
805 context.set_end("2021-01-01").unwrap();
806
807 assert!(Attachment::get_total_attachment_bytes(&connection, &context).is_ok());
808 }
809
810 #[test]
811 fn can_get_attachment_bytes_contact_filter() {
812 let db_path = current_dir()
813 .unwrap()
814 .parent()
815 .unwrap()
816 .join("imessage-database/test_data/db/test.db");
817 let connection = get_connection(&db_path).unwrap();
818
819 let mut context = QueryContext::default();
820 context.set_selected_chat_ids(BTreeSet::from([1, 2, 3]));
821 context.set_selected_handle_ids(BTreeSet::from([1, 2, 3]));
822
823 assert!(Attachment::get_total_attachment_bytes(&connection, &context).is_ok());
824 }
825
826 #[test]
827 fn can_get_attachment_bytes_contact_date_filter() {
828 let db_path = current_dir()
829 .unwrap()
830 .parent()
831 .unwrap()
832 .join("imessage-database/test_data/db/test.db");
833 let connection = get_connection(&db_path).unwrap();
834
835 let mut context = QueryContext::default();
836 context.set_start("2020-01-01").unwrap();
837 context.set_end("2021-01-01").unwrap();
838 context.set_selected_chat_ids(BTreeSet::from([1, 2, 3]));
839 context.set_selected_handle_ids(BTreeSet::from([1, 2, 3]));
840
841 assert!(Attachment::get_total_attachment_bytes(&connection, &context).is_ok());
842 }
843
844 #[test]
845 fn can_get_file_size_bytes() {
846 let attachment = sample_attachment();
847
848 assert_eq!(attachment.file_size(), String::from("100.00 B"));
849 }
850
851 #[test]
852 fn can_get_file_size_kb() {
853 let mut attachment = sample_attachment();
854 attachment.total_bytes = 2300;
855
856 assert_eq!(attachment.file_size(), String::from("2.25 KB"));
857 }
858
859 #[test]
860 fn can_get_file_size_mb() {
861 let mut attachment = sample_attachment();
862 attachment.total_bytes = 5612000;
863
864 assert_eq!(attachment.file_size(), String::from("5.35 MB"));
865 }
866
867 #[test]
868 fn can_get_file_size_gb() {
869 let mut attachment: Attachment = sample_attachment();
870 attachment.total_bytes = 9234712394;
871
872 assert_eq!(attachment.file_size(), String::from("8.60 GB"));
873 }
874
875 #[test]
876 fn can_get_file_size_cap() {
877 let mut attachment: Attachment = sample_attachment();
878 attachment.total_bytes = i64::MAX;
879
880 assert_eq!(attachment.file_size(), String::from("8388608.00 TB"));
881 }
882}