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}