tests/test_message_history_reader.py — Message History Reader for iMessage and SMS

#!/usr/bin/env python3
# Synthetic tests for Message History Reader for iMessage and SMS.
# By Zero Width Lab (zerowidthlab.com)
from __future__ import annotations

import importlib.util
import json
import sqlite3
import subprocess
import sys
import tempfile
import unittest
from pathlib import Path
from unittest.mock import patch

__author__ = "Zero Width Lab"
SCRIPT = Path(__file__).resolve().parents[1] / "scripts" / "message_history_reader.py"
SPEC = importlib.util.spec_from_file_location("message_history_reader", SCRIPT)
assert SPEC and SPEC.loader
reader = importlib.util.module_from_spec(SPEC)
SPEC.loader.exec_module(reader)


def typedstream(text: str) -> bytes:
    payload = text.encode("utf-8")
    return b"\x04\x0bstreamtyped\x81\xe8\x03NSString\x84" + len(payload).to_bytes(4, "little") + payload


def make_database(path: Path) -> None:
    connection = sqlite3.connect(path)
    connection.executescript("""
        CREATE TABLE chat (ROWID INTEGER PRIMARY KEY, chat_identifier TEXT);
        CREATE TABLE handle (ROWID INTEGER PRIMARY KEY, id TEXT);
        CREATE TABLE message (ROWID INTEGER PRIMARY KEY, text TEXT, attributedBody BLOB, date INTEGER, is_from_me INTEGER, handle_id INTEGER);
        CREATE TABLE chat_message_join (chat_id INTEGER, message_id INTEGER);
        CREATE TABLE chat_handle_join (chat_id INTEGER, handle_id INTEGER);
        INSERT INTO chat VALUES (1,'test-chat');
        INSERT INTO handle VALUES (1,'test-handle');
        INSERT INTO chat_handle_join VALUES (1,1);
        INSERT INTO message VALUES (1,'hello',NULL,700000000,0,1);
        INSERT INTO chat_message_join VALUES (1,1);
    """)
    connection.execute("INSERT INTO message VALUES (2,NULL,?,700000001000000000,1,1)", (typedstream("second line\nspoof"),))
    connection.execute("INSERT INTO chat_message_join VALUES (1,2)")
    connection.commit()
    connection.close()


class ReaderTests(unittest.TestCase):
    def setUp(self) -> None:
        self.temporary = tempfile.TemporaryDirectory()
        self.addCleanup(self.temporary.cleanup)
        self.root = Path(self.temporary.name).resolve()
        self.db = self.root / "chat.db"
        make_database(self.db)

    def cli(self, *arguments: str) -> subprocess.CompletedProcess[str]:
        return subprocess.run([sys.executable, str(SCRIPT), "--db", str(self.db), *arguments],
                              text=True, capture_output=True, check=False)

    def test_chats_and_search(self) -> None:
        with patch.object(reader.sys, "platform", "darwin"):
            with reader.snapshot(self.db) as connection:
                self.assertEqual(reader.chats(connection, 10)[0]["message_count"], 2)
                rows, skipped = reader.messages(connection, term="second", handle="test-handle")
        self.assertEqual(skipped, [])
        self.assertEqual(rows[0]["text"], "second line\nspoof")

    def test_group_chat_handle_and_sender_filters(self) -> None:
        connection = sqlite3.connect(self.db)
        connection.executescript("""
            INSERT INTO handle VALUES (2,'other-handle');
            INSERT INTO chat_handle_join VALUES (1,2);
            INSERT INTO message VALUES (3,'shared phrase',NULL,700000002,0,1);
            INSERT INTO message VALUES (4,'shared phrase',NULL,700000003,0,2);
            INSERT INTO message VALUES (5,'shared phrase',NULL,700000004,1,2);
            INSERT INTO chat_message_join VALUES (1,3);
            INSERT INTO chat_message_join VALUES (1,4);
            INSERT INTO chat_message_join VALUES (1,5);
        """)
        connection.commit()
        connection.close()
        with reader.snapshot(self.db) as snapshot:
            chat_rows, _ = reader.messages(snapshot, term="shared phrase", handle="test-handle")
            sender_rows, _ = reader.messages(snapshot, term="shared phrase", sender="test-handle")
            other_rows, _ = reader.messages(snapshot, term="shared phrase", sender="other-handle")
            combined_rows, _ = reader.messages(snapshot, term="shared phrase", handle="test-handle", sender="other-handle")
        self.assertEqual([row["message_id"] for row in chat_rows], [3, 4, 5])
        self.assertEqual([row["message_id"] for row in sender_rows], [3])
        self.assertEqual([row["message_id"] for row in other_rows], [4])
        self.assertEqual([row["message_id"] for row in combined_rows], [4])

    def test_limit_stops_before_later_undecodable_body(self) -> None:
        connection = sqlite3.connect(self.db)
        connection.execute("UPDATE message SET attributedBody=? WHERE ROWID=2", (b"bad body",))
        connection.commit()
        connection.close()
        with reader.snapshot(self.db) as snapshot:
            rows, skipped = reader.messages(snapshot, term="hello", limit=1)
        self.assertEqual([row["message_id"] for row in rows], [1])
        self.assertEqual(skipped, [])

    def test_timestamp_units_and_escaping(self) -> None:
        self.assertEqual(reader.iso_date(700000000), reader.iso_date(700000000000000000))
        self.assertIn("\\u000a", reader.display_text("line\nspoof"))
        self.assertEqual(reader.extract_attributed_body(typedstream("hello")), "hello")

    def test_undecodable_is_incomplete(self) -> None:
        connection = sqlite3.connect(self.db)
        connection.execute("UPDATE message SET attributedBody=? WHERE ROWID=2", (b"bad body",))
        connection.commit()
        connection.close()
        with patch.object(reader.sys, "platform", "darwin"):
            self.assertEqual(reader.main(["--db", str(self.db), "search", "--text", "hello"]), 2)
            self.assertEqual(reader.main(["--db", str(self.db), "--allow-partial", "search", "--text", "hello"]), 0)

    def test_export_gate_and_create_only(self) -> None:
        with patch.object(reader.sys, "platform", "darwin"):
            out = self.root / "export"
            with patch("builtins.print") as printing:
                self.assertEqual(reader.main(["--db", str(self.db), "--json", "export", "--chat", "1", "--out", str(out), "--dry-run"]), 0)
            preview = json.loads(printing.call_args.args[0])
            self.assertFalse(out.exists())
            with self.assertRaises(reader.ReaderError):
                reader.main(["--db", str(self.db), "export", "--chat", "1", "--out", str(out)])
            self.assertEqual(reader.main(["--db", str(self.db), "export", "--chat", "1", "--out", str(out), "--approve", preview["approval_digest"]]), 0)
            self.assertEqual(len((out / "messages.jsonl").read_text().splitlines()), 2)
            self.assertIn("\n    | spoof\n", (out / "transcript.txt").read_text())
            with self.assertRaises(reader.ReaderError):
                reader.create_export(out, {})

    def test_digest_changes_with_content_and_destination(self) -> None:
        first = reader.export_digest(1, self.root / "a", {"messages.jsonl": b"one"})
        self.assertNotEqual(first, reader.export_digest(1, self.root / "a", {"messages.jsonl": b"two"}))
        self.assertNotEqual(first, reader.export_digest(1, self.root / "b", {"messages.jsonl": b"one"}))

    def test_refuses_source_symlink(self) -> None:
        link = self.root / "link.db"
        link.symlink_to(self.db)
        with self.assertRaises(reader.ReaderError):
            with reader.snapshot(link):
                pass

    def test_refuses_export_output_symlink(self) -> None:
        actual = self.root / "actual"
        actual.mkdir()
        link = self.root / "export-link"
        link.symlink_to(actual, target_is_directory=True)
        with patch.object(reader.sys, "platform", "darwin"):
            with self.assertRaisesRegex(reader.ReaderError, "not a symlink"):
                reader.main(["--db", str(self.db), "export", "--chat", "1", "--out", str(link), "--dry-run"])
        self.assertEqual(list(actual.iterdir()), [])

    def test_refused_sidecars_mark_results_incomplete(self) -> None:
        (self.root / "chat.db-wal").symlink_to(self.db)
        (self.root / "chat.db-shm").mkdir()
        skipped: list[str] = []
        with reader.snapshot(self.db, skipped):
            pass
        self.assertEqual(len(skipped), 2)
        self.assertIn("chat.db-wal", skipped[0])
        self.assertIn("chat.db-shm", skipped[1])
        with patch.object(reader.sys, "platform", "darwin"):
            with patch("builtins.print") as printing:
                self.assertEqual(reader.main(["--db", str(self.db), "--json", "chats"]), 2)
        result = json.loads(next(call.args[0] for call in printing.call_args_list if call.kwargs.get("file") is None))
        self.assertTrue(result["incomplete"])
        self.assertEqual(len(result["skipped"]), 2)

    def test_missing_database_mentions_full_disk_access(self) -> None:
        with self.assertRaisesRegex(reader.ReaderError, "Full Disk Access"):
            with reader.snapshot(self.root / "missing.db"):
                pass

    def test_parent_symlink_swap_refused_before_write(self) -> None:
        parent = self.root / "parent"
        parent.mkdir()
        out = parent / "export"
        plan = reader.output_plan(out)
        original = self.root / "original"
        parent.rename(original)
        redirected = self.root / "redirected"
        redirected.mkdir()
        parent.symlink_to(redirected, target_is_directory=True)
        with self.assertRaises(OSError):
            reader.create_export(out, {"messages.jsonl": b"private"}, plan)
        self.assertEqual(list(redirected.iterdir()), [])
        self.assertEqual(list(original.iterdir()), [])

    def test_existing_directory_identity_is_bound(self) -> None:
        out = self.root / "export"
        out.mkdir()
        plan = reader.output_plan(out)
        digest = reader.export_digest(1, out, {}, plan)
        out.rename(self.root / "old")
        out.mkdir()
        self.assertNotEqual(digest, reader.export_digest(1, out, {}))
        with self.assertRaisesRegex(reader.ReaderError, "changed"):
            reader.create_export(out, {}, plan)

    def test_no_creation_above_output(self) -> None:
        with self.assertRaises(reader.ReaderError):
            reader.output_plan(self.root / "missing" / "export")
        self.assertFalse((self.root / "missing").exists())
        self.assertTrue(reader.output_plan(self.root / "export")["create_directory"])

    def test_traversal_rejected_before_reads(self) -> None:
        for suffix in ("../outside.db", "%2e%2e/outside.db", "folder\\outside.db"):
            with self.subTest(path=suffix), patch.object(reader.os, "open") as opening:
                with self.assertRaises(reader.ReaderError):
                    with reader.snapshot(self.root / suffix):
                        pass
                opening.assert_not_called()
        parent = self.root / "link"
        parent.symlink_to(self.root, target_is_directory=True)
        with self.assertRaises(OSError):
            with reader.snapshot(parent / "chat.db"):
                pass

    def test_unsupported_platform_before_io(self) -> None:
        with patch.object(reader.sys, "platform", "win32"), patch.object(reader, "snapshot") as copying:
            with self.assertRaisesRegex(reader.ReaderError, "macOS"):
                reader.main(["chats"])
            copying.assert_not_called()
        with patch.object(reader.sys, "platform", "darwin"), patch.object(reader.os, "supports_dir_fd", set()):
            with self.assertRaisesRegex(reader.ReaderError, "directory-descriptor"):
                reader.main(["chats"])

    def test_input_limit_and_malformed_file(self) -> None:
        with patch.object(reader, "MAX_SOURCE_BYTES", 1):
            with self.assertRaisesRegex(reader.ReaderError, "128 MiB"):
                with reader.snapshot(self.db):
                    pass
        bad = self.root / "bad.db"
        bad.write_bytes(b"not a database")
        with self.assertRaises(sqlite3.DatabaseError):
            with reader.snapshot(bad):
                pass
        connection = sqlite3.connect(self.db)
        connection.execute("UPDATE message SET text=? WHERE ROWID=1", (b"not text",))
        connection.commit()
        connection.close()
        with reader.snapshot(self.db) as copied:
            with self.assertRaisesRegex(reader.ReaderError, "text must be"):
                reader.messages(copied, term="anything")

    def test_malformed_chat_date_and_identifier(self) -> None:
        connection = sqlite3.connect(self.db)
        connection.execute("UPDATE message SET date='invalid'")
        connection.commit()
        connection.close()
        with reader.snapshot(self.db) as copied:
            with self.assertRaisesRegex(reader.ReaderError, "last date"):
                reader.chats(copied, 10)
        connection = sqlite3.connect(self.db)
        connection.execute("UPDATE chat SET chat_identifier=?", (b"invalid",))
        connection.commit()
        connection.close()
        with reader.snapshot(self.db) as copied:
            with self.assertRaisesRegex(reader.ReaderError, "identifier"):
                reader.chats(copied, 10)

    def test_human_preview_discloses_creation(self) -> None:
        with patch.object(reader.sys, "platform", "darwin"), patch("builtins.print") as printing:
            reader.main(["--db", str(self.db), "export", "--chat", "1", "--out", str(self.root / "export"), "--dry-run"])
        lines = [call.args[0] for call in printing.call_args_list]
        self.assertTrue(any("Directory action: create" in line for line in lines))
        self.assertTrue(any("Parent identity:" in line for line in lines))
        self.assertFalse((self.root / "export").exists())

    def test_cli_subprocess_and_help(self) -> None:
        result = self.cli("--json", "chats", "--limit", "1")
        self.assertEqual(result.returncode, 0, result.stderr)
        self.assertEqual(json.loads(result.stdout)["results"][0]["message_count"], 2)
        help_result = subprocess.run([sys.executable, str(SCRIPT), "--help"],
                                     text=True, capture_output=True, check=False)
        self.assertEqual(help_result.returncode, 0)
        self.assertTrue(help_result.stdout.rstrip().endswith(reader.SIGNATURE))
        search_help = self.cli("search", "--help")
        self.assertEqual(search_help.returncode, 0)
        self.assertIn("chats that include this handle", search_help.stdout)
        self.assertIn("messages sent by this handle", search_help.stdout)


if __name__ == "__main__":
    unittest.main()

License: MIT (read it).

Back to the listing