wz_reader 0.0.14

A wz file reader to resolve wz file with thread safe
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128

use scroll::{Pread, LE};

use super::snow2_decryptor::Snow2Decryptor;
use crate::reader::Error;

pub(crate) struct Snow2Reader<'a> {
    data: &'a [u8],
    pub offset: usize,
    decryptor: Snow2Decryptor,
    pub buffer: [u8; 4],
    pub buffer_len: usize,
}

impl<'a> Snow2Reader<'a> {
    pub fn new(data: &'a [u8], snow_key: [u8; 16]) -> Self {
        Self {
            data,
            offset: 0,
            decryptor: Snow2Decryptor::new(snow_key),
            buffer: [0_u8; 4],
            buffer_len: 0,
        }
    }
    #[allow(dead_code)]
    pub fn read_u32(&mut self) -> Result<u32, Error> {
        let decrypted = self
            .decryptor
            .decrypt_block(&self.data.pread_with::<u32>(self.offset, LE)?);
        self.offset += 4;
        if self.buffer_len > 0 {
            let decrypted_bytes = decrypted.to_le_bytes();
            let mut buffer = [0_u8; 4];
            buffer[..self.buffer_len].copy_from_slice(&self.buffer[..self.buffer_len]);
            buffer[self.buffer_len..].copy_from_slice(&decrypted_bytes[..4 - self.buffer_len]);
            let result = buffer.pread_with::<u32>(0, LE)?;
            self.buffer[..(4 - self.buffer_len)]
                .copy_from_slice(&decrypted_bytes[self.buffer_len..]);
            Ok(result)
        } else {
            Ok(decrypted)
        }
    }
    pub fn read_i32(&mut self) -> Result<i32, Error> {
        let decrypted = self
            .decryptor
            .decrypt_block(&self.data.pread_with::<u32>(self.offset, LE)?);
        self.offset += 4;

        let decrypted_bytes = decrypted.to_le_bytes();

        if self.buffer_len > 0 {
            // merge buffer and decrypted_bytes then read u32
            let mut buffer = [0_u8; 4];
            buffer[..self.buffer_len].copy_from_slice(&self.buffer[..self.buffer_len]);
            buffer[self.buffer_len..].copy_from_slice(&decrypted_bytes[..4 - self.buffer_len]);
            let result = buffer.pread_with::<i32>(0, LE)?;
            // deal with remaining bytes
            self.buffer[..(4 - self.buffer_len)]
                .copy_from_slice(&decrypted_bytes[self.buffer_len..]);

            Ok(result)
        } else {
            Ok(i32::from_le_bytes(decrypted_bytes))
        }
    }
    #[inline]
    pub fn read_utf16_string(&mut self, len: usize) -> Result<String, Error> {
        let string_vec = self.read_bytes(len)?;
        let utf16_vec = string_vec
            .chunks_exact(2)
            .map(|u| u16::from_le_bytes([u[0], u[1]]))
            .collect::<Vec<_>>();

        String::from_utf16(&utf16_vec).map_err(Error::from)
    }
    #[inline]
    pub fn read_bytes(&mut self, len: usize) -> Result<Vec<u8>, Error> {
        let mut vec = vec![0_u8; len];

        self.write_bytes_to(&mut vec, len)?;

        Ok(vec)
    }
    pub fn write_bytes_to(&mut self, dest_buffer: &mut [u8], len: usize) -> Result<(), Error> {
        let mut remaining_len = len as i32;

        /* first: deal remain buffer from previous */
        if self.buffer_len > 0 {
            if len < self.buffer_len {
                dest_buffer[0..len].copy_from_slice(&self.buffer[0..len]);
                self.buffer_len -= len;

                // move the remaining bytes to the beginning of the buffer
                let mut temp_buffer = [0_u8; 4];
                temp_buffer.copy_from_slice(&self.buffer[len..]);

                self.buffer.copy_from_slice(&temp_buffer);

                return Ok(());
            }
            dest_buffer[0..self.buffer_len].copy_from_slice(&self.buffer[0..self.buffer_len]);
            // self.buffer.fill(0);
            remaining_len -= self.buffer_len as i32;
            self.buffer_len = 0;
        }

        while remaining_len > 0 {
            let decrypted = self
                .decryptor
                .decrypt_block(&self.data.pread_with::<u32>(self.offset, LE)?);
            let bytes = decrypted.to_le_bytes();
            let start = len - remaining_len as usize;
            let dest_range = start..std::cmp::min(start + 4, len);
            let data_len = dest_range.len();
            dest_buffer[dest_range].copy_from_slice(&bytes[..data_len]);
            self.offset += 4;
            remaining_len -= 4;

            if remaining_len < 0 {
                let remaining_bytes = &bytes[data_len..];
                self.buffer_len = remaining_bytes.len();
                self.buffer[..self.buffer_len].copy_from_slice(remaining_bytes);
            }
        }

        Ok(())
    }
}