msg_tool\utils/
rc4.rs

1use std::io::{Read, Write};
2
3pub struct Rc4 {
4    state: [u8; 256],
5    i: u8,
6    j: u8,
7}
8
9impl Rc4 {
10    pub fn new(key: &[u8]) -> Self {
11        let mut state = [0u8; 256];
12        for i in 0..256 {
13            state[i] = i as u8;
14        }
15
16        let mut j: u8 = 0;
17        for i in 0..256 {
18            j = j.wrapping_add(state[i]).wrapping_add(key[i % key.len()]);
19            state.swap(i, j as usize);
20        }
21
22        Rc4 { state, i: 0, j: 0 }
23    }
24
25    pub fn next_byte(&mut self) -> u8 {
26        self.i = self.i.wrapping_add(1);
27        self.j = self.j.wrapping_add(self.state[self.i as usize]);
28        self.state.swap(self.i as usize, self.j as usize);
29        let k = self.state
30            [(self.state[self.i as usize].wrapping_add(self.state[self.j as usize])) as usize];
31        k
32    }
33
34    pub fn skip_bytes(&mut self, n: usize) {
35        for _ in 0..n {
36            self.next_byte();
37        }
38    }
39
40    pub fn generate_block(&mut self, len: usize) -> Vec<u8> {
41        (0..len).map(|_| self.next_byte()).collect()
42    }
43
44    pub fn process_block(&mut self, data: &mut [u8]) {
45        for byte in data.iter_mut() {
46            *byte ^= self.next_byte();
47        }
48    }
49}
50
51pub struct Rc4Stream<T> {
52    inner: T,
53    rc4: Rc4,
54}
55
56impl<T> Rc4Stream<T> {
57    pub fn new(inner: T, rc4: Rc4) -> Self {
58        Rc4Stream { inner, rc4 }
59    }
60
61    pub fn new_with_key(inner: T, key: &[u8]) -> Self {
62        Rc4Stream {
63            inner,
64            rc4: Rc4::new(key),
65        }
66    }
67}
68
69impl<T: Read> Read for Rc4Stream<T> {
70    fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
71        let n = self.inner.read(buf)?;
72        self.rc4.process_block(&mut buf[..n]);
73        Ok(n)
74    }
75}
76
77impl<T: Write> Write for Rc4Stream<T> {
78    fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
79        let mut data = buf.to_vec();
80        self.rc4.process_block(&mut data);
81        self.inner.write(&data)
82    }
83
84    fn flush(&mut self) -> std::io::Result<()> {
85        self.inner.flush()
86    }
87}