imessage_database/message_types/
polls.rs

1/*!
2 These are the poll messages generated by the iOS 26 (or newer) Polls app.
3*/
4
5use std::collections::HashMap;
6
7use base64::{Engine as _, engine::general_purpose};
8use jzon::{self, JsonValue};
9use plist::Value;
10
11use crate::{
12    error::plist::PlistParseError,
13    util::plist::{get_string_from_nested_dict, parse_ns_keyed_archiver},
14};
15
16/// Type alias for Poll Option ID
17pub type PollOptionID = String;
18
19/// Represents a Poll option
20#[derive(Debug, PartialEq, Eq)]
21pub struct PollOption {
22    /// The text of the option
23    pub text: String,
24    /// The creator of the option
25    pub creator: String,
26    /// The votes for this option
27    pub votes: Vec<PollVote>,
28}
29
30impl PollOption {
31    fn from_json(json_data: &JsonValue) -> (PollOptionID, Self) {
32        let poll_id = json_data["optionIdentifier"]
33            .as_str()
34            .unwrap_or_default()
35            .to_string();
36        let creator = json_data["creatorHandle"]
37            .as_str()
38            .unwrap_or_default()
39            .to_string();
40        let text = json_data["text"].as_str().unwrap_or_default().to_string();
41
42        (
43            poll_id,
44            PollOption {
45                text,
46                creator,
47                votes: Vec::new(),
48            },
49        )
50    }
51}
52
53/// Represents a vote in a Poll
54#[derive(Debug, PartialEq, Eq)]
55pub struct PollVote {
56    /// The handle of the voter
57    pub voter: String,
58    /// The ID of the option being voted on
59    pub option_id: PollOptionID,
60}
61
62impl PollVote {
63    fn from_json(payload: &JsonValue) -> Result<Self, PlistParseError> {
64        let voter = payload["participantHandle"]
65            .as_str()
66            .ok_or(PlistParseError::MissingKey("participantHandle".to_string()))?
67            .to_string();
68
69        let option_id = payload["voteOptionIdentifier"]
70            .as_str()
71            .ok_or(PlistParseError::MissingKey(
72                "voteOptionIdentifier".to_string(),
73            ))?
74            .to_string();
75
76        Ok(PollVote { voter, option_id })
77    }
78}
79
80/// Represents a Poll message
81#[derive(Debug, PartialEq, Eq)]
82pub struct Poll {
83    /// Map of option ID to [`PollOption`]
84    pub options: HashMap<PollOptionID, PollOption>,
85    /// The order of the options as they were created
86    pub order: Vec<PollOptionID>,
87}
88
89impl Poll {
90    /// Parse a Poll from the given payload
91    pub fn from_payload(payload: &Value) -> Result<Self, PlistParseError> {
92        let parsed = parse_ns_keyed_archiver(payload)?;
93
94        let url = get_string_from_nested_dict(&parsed, "URL")
95            .ok_or(PlistParseError::MissingKey("URL".to_string()))?;
96
97        // Parse the JSON
98        let parsed_json = base64_url_to_json(url)?;
99
100        let mut options = HashMap::new();
101        let mut order = Vec::new();
102        let ordered_options = parsed_json["item"]["orderedPollOptions"].as_array().ok_or(
103            PlistParseError::MissingKey("orderedPollOptions".to_string()),
104        )?;
105
106        for option in ordered_options {
107            let (id, poll_option) = PollOption::from_json(option);
108            order.push(id.clone());
109            options.insert(id, poll_option);
110        }
111
112        Ok(Poll { options, order })
113    }
114
115    /// Count votes from a vote payload and update the poll options
116    pub fn count_votes(&mut self, payload: &Value) -> Result<(), PlistParseError> {
117        let parsed = parse_ns_keyed_archiver(payload)?;
118
119        let url = get_string_from_nested_dict(&parsed, "URL")
120            .ok_or(PlistParseError::MissingKey("URL".to_string()))?;
121
122        // Parse the JSON
123        let parsed_json = base64_url_to_json(url)?;
124
125        let votes = parsed_json["item"]["votes"]
126            .as_array()
127            .ok_or(PlistParseError::MissingKey("votes".to_string()))?;
128
129        for vote in votes {
130            let poll_vote = PollVote::from_json(vote)?;
131            if let Some(option) = self.options.get_mut(&poll_vote.option_id) {
132                option.votes.push(poll_vote);
133            }
134        }
135
136        Ok(())
137    }
138}
139
140fn base64_url_to_json(data: &str) -> Result<JsonValue, PlistParseError> {
141    // Strip the fixed prefix
142    let after_prefix = data
143        .strip_prefix("data:,")
144        .ok_or(PlistParseError::WrongMessageType)?;
145
146    // Extract the base64 part before the first "?" (if present)
147    let base64_part = after_prefix
148        .split_once('?')
149        .map_or(after_prefix, |(before, _)| before);
150
151    // Decode the base64 part
152    let bytes = String::from_utf8(
153        general_purpose::URL_SAFE
154            .decode(base64_part)
155            .map_err(|_| PlistParseError::WrongMessageType)?,
156    )
157    .map_err(|_| PlistParseError::WrongMessageType)?;
158
159    // Parse the JSON
160    jzon::parse(&bytes).map_err(|_| PlistParseError::WrongMessageType)
161}
162
163#[cfg(test)]
164mod tests {
165    use std::{env::current_dir, fs::File};
166
167    use plist::Value;
168
169    use crate::message_types::polls::Poll;
170
171    #[test]
172    fn test_parse_poll_creation() {
173        let plist_path = current_dir()
174            .unwrap()
175            .as_path()
176            .join("test_data/app_message/PollCreate.plist");
177        let plist_data = File::open(plist_path).unwrap();
178        let plist = Value::from_reader(plist_data).unwrap();
179        println!("{:#?}", plist);
180
181        let poll = Poll::from_payload(&plist).unwrap();
182
183        println!("{:#?}", poll);
184        assert_eq!(poll.options.len(), 4);
185    }
186
187    #[test]
188    fn test_parse_poll_votes() {
189        // Parse the poll first
190        let plist_path = current_dir()
191            .unwrap()
192            .as_path()
193            .join("test_data/app_message/PollCreate.plist");
194        let plist_data = File::open(plist_path).unwrap();
195        let plist = Value::from_reader(plist_data).unwrap();
196
197        let mut poll = Poll::from_payload(&plist).unwrap();
198
199        let plist_path = current_dir()
200            .unwrap()
201            .as_path()
202            .join("test_data/app_message/PollVote.plist");
203        let plist_data = File::open(plist_path).unwrap();
204        let plist = Value::from_reader(plist_data).unwrap();
205
206        poll.count_votes(&plist).unwrap();
207
208        let plist_path = current_dir()
209            .unwrap()
210            .as_path()
211            .join("test_data/app_message/PollVote2.plist");
212        let plist_data = File::open(plist_path).unwrap();
213        let plist = Value::from_reader(plist_data).unwrap();
214
215        poll.count_votes(&plist).unwrap();
216
217        println!("{:#?}", poll);
218    }
219
220    #[test]
221    fn test_parse_poll_vote_removed() {
222        // Parse the poll first
223        let plist_path = current_dir()
224            .unwrap()
225            .as_path()
226            .join("test_data/app_message/PollCreate.plist");
227        let plist_data = File::open(plist_path).unwrap();
228        let plist = Value::from_reader(plist_data).unwrap();
229
230        let mut poll = Poll::from_payload(&plist).unwrap();
231
232        let plist_path = current_dir()
233            .unwrap()
234            .as_path()
235            .join("test_data/app_message/PollRemovedVote.plist");
236        let plist_data = File::open(plist_path).unwrap();
237        let plist = Value::from_reader(plist_data).unwrap();
238
239        poll.count_votes(&plist).unwrap();
240        println!("{:#?}", poll);
241    }
242}
imessage_database/message_types/polls.rs

imessage_database/message_types/
polls.rs
