msg_tool\utils/
xored_stream.rs

1use std::io::{Read, Seek, Write};
2
3pub struct XoredStream<T> {
4    reader: T,
5    key: u8,
6}
7
8impl<T> XoredStream<T> {
9    pub fn new(reader: T, key: u8) -> Self {
10        XoredStream { reader, key }
11    }
12}
13
14impl<T: Read> Read for XoredStream<T> {
15    fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
16        let read_bytes = self.reader.read(buf)?;
17        for byte in &mut buf[..read_bytes] {
18            *byte ^= self.key;
19        }
20        Ok(read_bytes)
21    }
22}
23
24impl<T: Seek> Seek for XoredStream<T> {
25    fn seek(&mut self, pos: std::io::SeekFrom) -> std::io::Result<u64> {
26        self.reader.seek(pos)
27    }
28
29    fn rewind(&mut self) -> std::io::Result<()> {
30        self.reader.rewind()
31    }
32
33    fn stream_position(&mut self) -> std::io::Result<u64> {
34        self.reader.stream_position()
35    }
36}
37
38impl<T: Write> Write for XoredStream<T> {
39    fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
40        let mut encrypted_buf = buf.to_vec();
41        for byte in &mut encrypted_buf {
42            *byte ^= self.key;
43        }
44        self.reader.write(&encrypted_buf)
45    }
46
47    fn flush(&mut self) -> std::io::Result<()> {
48        self.reader.flush()
49    }
50}
51
52impl<T: std::fmt::Debug> std::fmt::Debug for XoredStream<T> {
53    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
54        f.debug_struct("XoredStream")
55            .field("reader", &self.reader)
56            .field("key", &self.key)
57            .finish()
58    }
59}
60
61/// A stream that XORs data with a repeating key based on the current position.
62pub struct XoredKeyStream<T> {
63    inner: T,
64    key: Vec<u8>,
65    base_position: u64,
66}
67
68impl<T> XoredKeyStream<T> {
69    pub fn new(inner: T, key: Vec<u8>, base_position: u64) -> Self {
70        XoredKeyStream {
71            inner,
72            key,
73            base_position,
74        }
75    }
76}
77
78impl<T: Read + Seek> Read for XoredKeyStream<T> {
79    fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
80        let key_len = self.key.len();
81        let start_pos =
82            ((self.inner.stream_position()? + self.base_position) % (key_len as u64)) as usize;
83        let readed = self.inner.read(buf)?;
84        for i in 0..readed {
85            buf[i] ^= self.key[(start_pos + i) % key_len];
86        }
87        Ok(readed)
88    }
89}
90
91impl<T: Seek> Seek for XoredKeyStream<T> {
92    fn seek(&mut self, pos: std::io::SeekFrom) -> std::io::Result<u64> {
93        self.inner.seek(pos)
94    }
95
96    fn rewind(&mut self) -> std::io::Result<()> {
97        self.inner.rewind()
98    }
99
100    fn stream_position(&mut self) -> std::io::Result<u64> {
101        self.inner.stream_position()
102    }
103}
104
105impl<T: Write + Seek> Write for XoredKeyStream<T> {
106    fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
107        let key_len = self.key.len();
108        let start_pos =
109            ((self.inner.stream_position()? + self.base_position) % (key_len as u64)) as usize;
110        let mut encrypted_buf = buf.to_vec();
111        for i in 0..buf.len() {
112            encrypted_buf[i] ^= self.key[(start_pos + i) % key_len];
113        }
114        self.inner.write(&encrypted_buf)
115    }
116
117    fn flush(&mut self) -> std::io::Result<()> {
118        self.inner.flush()
119    }
120}
121
122impl<T: std::fmt::Debug> std::fmt::Debug for XoredKeyStream<T> {
123    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
124        f.debug_struct("XoredKeyStream")
125            .field("inner", &self.inner)
126            .field("base_position", &self.base_position)
127            .finish()
128    }
129}