1use crate::{Record, token};
4
5const TOKENIZED_SCRIPT_SUBTYPE: &str = ">>>AUTOIT SCRIPT<<<";
6const UNICODE_SCRIPT_SUBTYPE: &str = ">AUTOIT UNICODE SCRIPT<";
7const PLAIN_SCRIPT_SUBTYPE: &str = ">AUTOIT SCRIPT<";
8const AHK_SCRIPT_SUBTYPE: &str = ">AUTOHOTKEY SCRIPT<";
10const AHK_WITH_ICON_SUBTYPE: &str = ">AHK WITH ICON<";
11
12#[derive(Debug, Clone, PartialEq)]
14pub struct Script {
15 record_index: usize,
16 name: String,
17 kind: ScriptKind,
18 bytes: Vec<u8>,
19 text: Option<ScriptText>,
20 decode_error: Option<ScriptDecodeError>,
21 token_stream: Option<token::TokenStream>,
22 creation_time: u64,
23 last_write_time: u64,
24}
25
26impl Script {
27 #[must_use]
33 pub const fn record_index(&self) -> usize {
34 self.record_index
35 }
36
37 #[must_use]
43 pub fn name(&self) -> &str {
44 self.name.as_str()
45 }
46
47 #[must_use]
53 pub const fn kind(&self) -> ScriptKind {
54 self.kind
55 }
56
57 #[must_use]
64 pub fn bytes(&self) -> &[u8] {
65 self.bytes.as_slice()
66 }
67
68 #[must_use]
75 pub const fn text(&self) -> Option<&ScriptText> {
76 self.text.as_ref()
77 }
78
79 #[must_use]
86 pub const fn decode_error(&self) -> Option<ScriptDecodeError> {
87 self.decode_error
88 }
89
90 #[must_use]
97 pub const fn token_stream(&self) -> Option<&token::TokenStream> {
98 self.token_stream.as_ref()
99 }
100
101 #[must_use]
108 pub fn source_text(&self) -> Option<&str> {
109 self.text.as_ref().map(ScriptText::text)
110 }
111
112 #[must_use]
118 pub const fn creation_time(&self) -> u64 {
119 self.creation_time
120 }
121
122 #[must_use]
129 pub const fn last_write_time(&self) -> u64 {
130 self.last_write_time
131 }
132}
133
134#[derive(Debug, Clone, Copy, PartialEq, Eq)]
136pub enum ScriptKind {
137 Tokenized,
140 UnicodeText,
142 PlainText,
144}
145
146#[derive(Debug, Clone, Copy, PartialEq, Eq)]
148pub enum ScriptDecodeError {
149 OddUtf16ByteLength,
151 Token(token::TokenError),
153}
154
155#[derive(Debug, Clone, PartialEq, Eq)]
157pub struct ScriptText {
158 encoding: ScriptTextEncoding,
159 text: String,
160}
161
162impl ScriptText {
163 #[must_use]
169 pub const fn encoding(&self) -> ScriptTextEncoding {
170 self.encoding
171 }
172
173 #[must_use]
179 pub fn text(&self) -> &str {
180 self.text.as_str()
181 }
182}
183
184#[derive(Debug, Clone, Copy, PartialEq, Eq)]
186pub enum ScriptTextEncoding {
187 Utf8Lossy,
189 Utf16LeLossy,
191 TokenRender,
193}
194
195#[must_use]
206pub fn recover_scripts(records: &[Record]) -> Vec<Script> {
207 records.iter().filter_map(recover_script).collect()
208}
209
210fn recover_script(record: &Record) -> Option<Script> {
227 let kind = match record.subtype() {
228 TOKENIZED_SCRIPT_SUBTYPE => ScriptKind::Tokenized,
229 UNICODE_SCRIPT_SUBTYPE => ScriptKind::UnicodeText,
230 PLAIN_SCRIPT_SUBTYPE | AHK_SCRIPT_SUBTYPE | AHK_WITH_ICON_SUBTYPE => ScriptKind::PlainText,
231 _ => return None,
232 };
233 let bytes = record.payload_data().to_vec();
234 let mut token_stream = None;
235 let (text, decode_error) = match kind {
236 ScriptKind::Tokenized => match token::parse(bytes.as_slice()) {
237 Ok(stream) => (
238 {
239 let rendered = stream.render_source();
240 token_stream = Some(stream);
241 Some(ScriptText {
242 encoding: ScriptTextEncoding::TokenRender,
243 text: rendered,
244 })
245 },
246 None,
247 ),
248 Err(err) => (None, Some(ScriptDecodeError::Token(err))),
249 },
250 ScriptKind::UnicodeText => match decode_utf16_lossy(bytes.as_slice()) {
251 Ok(text) => (
252 Some(ScriptText {
253 encoding: ScriptTextEncoding::Utf16LeLossy,
254 text,
255 }),
256 None,
257 ),
258 Err(err) => (None, Some(err)),
259 },
260 ScriptKind::PlainText => (
261 Some(ScriptText {
262 encoding: ScriptTextEncoding::Utf8Lossy,
263 text: String::from_utf8_lossy(bytes.as_slice()).into_owned(),
264 }),
265 None,
266 ),
267 };
268 Some(Script {
269 record_index: record.index(),
270 name: record.name().to_string(),
271 kind,
272 bytes,
273 text,
274 decode_error,
275 token_stream,
276 creation_time: record.creation_time(),
277 last_write_time: record.last_write_time(),
278 })
279}
280
281fn decode_utf16_lossy(data: &[u8]) -> Result<String, ScriptDecodeError> {
296 let chunks = data.chunks_exact(2);
297 if !chunks.remainder().is_empty() {
298 return Err(ScriptDecodeError::OddUtf16ByteLength);
299 }
300 let code_units: Result<Vec<u16>, ScriptDecodeError> = chunks
301 .map(|chunk| {
302 let bytes: [u8; 2] = chunk
303 .try_into()
304 .map_err(|_err| ScriptDecodeError::OddUtf16ByteLength)?;
305 Ok(u16::from_le_bytes(bytes))
306 })
307 .collect();
308 Ok(String::from_utf16_lossy(code_units?.as_slice()))
309}
310
311#[cfg(test)]
312mod tests {
313 use super::*;
314 use crate::au3::{DecodedString, DecompressionStatus, RecordTestParts};
315
316 #[test]
317 fn recovers_plain_script_text() -> Result<(), String> {
318 let record = test_record(">AUTOIT SCRIPT<", "main.au3", b"MsgBox(0, \"x\", \"y\")")?;
319 let scripts = recover_scripts(&[record]);
320 let script = scripts
321 .first()
322 .ok_or_else(|| "missing script".to_string())?;
323
324 check_eq(script.record_index(), 7, "record index")?;
325 check_eq(script.name(), "main.au3", "name")?;
326 check_eq(script.kind(), ScriptKind::PlainText, "kind")?;
327 check_eq(script.creation_time(), 13, "creation time")?;
328 check_eq(script.last_write_time(), 17, "last-write time")?;
329 check_eq(
330 script.source_text(),
331 Some("MsgBox(0, \"x\", \"y\")"),
332 "text",
333 )
334 }
335
336 #[test]
337 fn recovers_utf16_script_text() -> Result<(), String> {
338 let mut payload = Vec::new();
339 for unit in "MsgBox(0, \"x\", \"y\")".encode_utf16() {
340 payload.extend_from_slice(&unit.to_le_bytes());
341 }
342 let record = test_record(">AUTOIT UNICODE SCRIPT<", "unicode.au3", payload.as_slice())?;
343 let scripts = recover_scripts(&[record]);
344 let script = scripts
345 .first()
346 .ok_or_else(|| "missing script".to_string())?;
347
348 check_eq(script.kind(), ScriptKind::UnicodeText, "kind")?;
349 check_eq(script.decode_error(), None, "decode error")?;
350 check_eq(
351 script.source_text(),
352 Some("MsgBox(0, \"x\", \"y\")"),
353 "text",
354 )
355 }
356
357 #[test]
358 fn preserves_tokenized_script_bytes_without_text() -> Result<(), String> {
359 let tokenized = tokenized_assignment()?;
360 let record = test_record(">>>AUTOIT SCRIPT<<<", "tokenized.au3", tokenized.as_slice())?;
361 let scripts = recover_scripts(&[record]);
362 let script = scripts
363 .first()
364 .ok_or_else(|| "missing script".to_string())?;
365
366 check_eq(script.kind(), ScriptKind::Tokenized, "kind")?;
367 check_eq(script.bytes(), tokenized.as_slice(), "bytes")?;
368 check_eq(script.source_text(), Some("$x = 1\r\n"), "text")?;
369 check_eq(script.token_stream().is_some(), true, "token stream")
370 }
371
372 #[test]
373 fn renders_tokenized_msgbox_script() -> Result<(), String> {
374 let tokenized = tokenized_msgbox()?;
375 let record = test_record(">>>AUTOIT SCRIPT<<<", "msgbox.au3", tokenized.as_slice())?;
376 let scripts = recover_scripts(&[record]);
377 let script = scripts
378 .first()
379 .ok_or_else(|| "missing script".to_string())?;
380
381 check_eq(script.kind(), ScriptKind::Tokenized, "kind")?;
382 check_eq(
383 script.source_text(),
384 Some("MsgBox(0, \"title\", \"text\")\r\n"),
385 "text",
386 )
387 }
388
389 #[test]
390 fn preserves_bad_utf16_script_bytes_with_decode_error() -> Result<(), String> {
391 let record = test_record(">AUTOIT UNICODE SCRIPT<", "bad.au3", b"\xff")?;
392 let scripts = recover_scripts(&[record]);
393 let script = scripts
394 .first()
395 .ok_or_else(|| "missing script".to_string())?;
396
397 check_eq(script.bytes(), b"\xff".as_slice(), "bytes")?;
398 check_eq(script.source_text(), None, "text")?;
399 check_eq(
400 script.decode_error(),
401 Some(ScriptDecodeError::OddUtf16ByteLength),
402 "decode error",
403 )
404 }
405
406 fn test_record(subtype: &str, name: &str, payload: &[u8]) -> Result<Record, String> {
407 let payload_len = u32::try_from(payload.len()).map_err(|err| err.to_string())?;
408 Ok(Record::from_parts_for_test(RecordTestParts {
409 index: 7,
410 offset: 11,
411 subtype: DecodedString::from_text_for_test(subtype),
412 name: DecodedString::from_text_for_test(name),
413 compressed: false,
414 compressed_size: payload_len,
415 uncompressed_size: payload_len,
416 checksum: 0,
417 checksum_valid: false,
418 creation_time: 13,
419 last_write_time: 17,
420 encrypted_data: payload.to_vec(),
421 decrypted_data: payload.to_vec(),
422 decompressed_data: None,
423 decompression_status: DecompressionStatus::NotCompressed,
424 profile: crate::RecordProfile {
425 encoding: crate::Encoding::Ea06,
426 encryption: crate::EncryptionProfile::Ea06Lame,
427 compression: crate::CompressionProfile::None,
428 },
429 }))
430 }
431
432 fn tokenized_assignment() -> Result<Vec<u8>, String> {
433 let mut data = Vec::new();
434 data.extend_from_slice(&1u32.to_le_bytes());
435 data.push(0x33);
436 append_xored_string(&mut data, "x")?;
437 data.push(0x41);
438 data.push(0x05);
439 data.extend_from_slice(&1u32.to_le_bytes());
440 data.push(0x7f);
441 Ok(data)
442 }
443
444 fn tokenized_msgbox() -> Result<Vec<u8>, String> {
445 let mut data = Vec::new();
446 data.extend_from_slice(&1u32.to_le_bytes());
447 data.push(0x01);
448 data.extend_from_slice(&248i32.to_le_bytes());
449 data.push(0x47);
450 data.push(0x05);
451 data.extend_from_slice(&0u32.to_le_bytes());
452 data.push(0x40);
453 data.push(0x36);
454 append_xored_string(&mut data, "title")?;
455 data.push(0x40);
456 data.push(0x36);
457 append_xored_string(&mut data, "text")?;
458 data.push(0x48);
459 data.push(0x7f);
460 Ok(data)
461 }
462
463 fn append_xored_string(out: &mut Vec<u8>, value: &str) -> Result<(), String> {
464 let units: Vec<u16> = value.encode_utf16().collect();
465 let key = u32::try_from(units.len()).map_err(|err| err.to_string())?;
466 out.extend_from_slice(&key.to_le_bytes());
467 let key16 = u16::try_from(key).map_err(|err| err.to_string())?;
468 for unit in units {
469 out.extend_from_slice(&(unit ^ key16).to_le_bytes());
470 }
471 Ok(())
472 }
473
474 fn check_eq<T>(actual: T, expected: T, context: &str) -> Result<(), String>
475 where
476 T: core::fmt::Debug + PartialEq,
477 {
478 if actual == expected {
479 Ok(())
480 } else {
481 Err(format!("{context}: got {actual:?}, expected {expected:?}"))
482 }
483 }
484}