1use crate::wire::stream::Cipher;
20use sha2::{Digest, Sha256};
21
22#[derive(Debug, Clone, Copy, PartialEq, Eq)]
24pub enum WireCryptPlugin {
25 Arc4,
27 ChaCha,
29 ChaCha64,
31}
32
33impl WireCryptPlugin {
34 pub fn name(self) -> &'static str {
36 match self {
37 WireCryptPlugin::Arc4 => "Arc4",
38 WireCryptPlugin::ChaCha => "ChaCha",
39 WireCryptPlugin::ChaCha64 => "ChaCha64",
40 }
41 }
42}
43
44#[derive(Clone)]
46pub struct Rc4 {
47 s: [u8; 256],
48 i: u8,
49 j: u8,
50}
51
52impl Rc4 {
53 pub fn new(key: &[u8]) -> Self {
55 assert!(!key.is_empty(), "RC4 key must be non-empty");
56 let mut s = [0u8; 256];
57 for (i, b) in s.iter_mut().enumerate() {
58 *b = i as u8;
59 }
60 let mut j: u8 = 0;
61 for i in 0..256 {
62 j = j.wrapping_add(s[i]).wrapping_add(key[i % key.len()]);
63 s.swap(i, j as usize);
64 }
65 Rc4 { s, i: 0, j: 0 }
66 }
67
68 #[inline]
69 fn next_byte(&mut self) -> u8 {
70 self.i = self.i.wrapping_add(1);
71 self.j = self.j.wrapping_add(self.s[self.i as usize]);
72 self.s.swap(self.i as usize, self.j as usize);
73 let idx = self.s[self.i as usize].wrapping_add(self.s[self.j as usize]);
74 self.s[idx as usize]
75 }
76}
77
78impl Cipher for Rc4 {
79 fn process(&mut self, data: &mut [u8]) {
80 for b in data.iter_mut() {
81 *b ^= self.next_byte();
82 }
83 }
84}
85
86#[derive(Clone)]
92pub struct ChaCha20 {
93 state: [u32; 16],
94 block: [u8; 64],
95 pos: usize,
96 wide_counter: bool,
98}
99
100const CHACHA_CONST: [u32; 4] = [0x6170_7865, 0x3320_646e, 0x7962_2d32, 0x6b20_6574];
101
102impl ChaCha20 {
103 pub fn new(key: &[u8], nonce: &[u8]) -> Self {
107 assert_eq!(key.len(), 32, "ChaCha20 key must be 32 bytes");
108 let mut state = [0u32; 16];
109 state[0..4].copy_from_slice(&CHACHA_CONST);
110 for i in 0..8 {
111 state[4 + i] = u32::from_le_bytes(key[i * 4..i * 4 + 4].try_into().unwrap());
112 }
113 let wide_counter = match nonce.len() {
114 12 => {
115 for i in 0..3 {
117 state[13 + i] = u32::from_le_bytes(nonce[i * 4..i * 4 + 4].try_into().unwrap());
118 }
119 false
120 }
121 8 => {
122 state[14] = u32::from_le_bytes(nonce[0..4].try_into().unwrap());
124 state[15] = u32::from_le_bytes(nonce[4..8].try_into().unwrap());
125 true
126 }
127 other => panic!("ChaCha nonce must be 12 or 8 bytes, got {other}"),
128 };
129 let mut c = ChaCha20 {
130 state,
131 block: [0; 64],
132 pos: 64,
133 wide_counter,
134 };
135 c.refill();
136 c
137 }
138
139 #[inline]
140 fn quarter_round(x: &mut [u32; 16], a: usize, b: usize, c: usize, d: usize) {
141 x[a] = x[a].wrapping_add(x[b]);
142 x[d] = (x[d] ^ x[a]).rotate_left(16);
143 x[c] = x[c].wrapping_add(x[d]);
144 x[b] = (x[b] ^ x[c]).rotate_left(12);
145 x[a] = x[a].wrapping_add(x[b]);
146 x[d] = (x[d] ^ x[a]).rotate_left(8);
147 x[c] = x[c].wrapping_add(x[d]);
148 x[b] = (x[b] ^ x[c]).rotate_left(7);
149 }
150
151 fn refill(&mut self) {
153 let mut x = self.state;
154 for _ in 0..10 {
155 Self::quarter_round(&mut x, 0, 4, 8, 12);
157 Self::quarter_round(&mut x, 1, 5, 9, 13);
158 Self::quarter_round(&mut x, 2, 6, 10, 14);
159 Self::quarter_round(&mut x, 3, 7, 11, 15);
160 Self::quarter_round(&mut x, 0, 5, 10, 15);
162 Self::quarter_round(&mut x, 1, 6, 11, 12);
163 Self::quarter_round(&mut x, 2, 7, 8, 13);
164 Self::quarter_round(&mut x, 3, 4, 9, 14);
165 }
166 for (i, w) in x.iter_mut().enumerate() {
167 *w = w.wrapping_add(self.state[i]);
168 self.block[i * 4..i * 4 + 4].copy_from_slice(&w.to_le_bytes());
169 }
170 self.pos = 0;
171 let (c0, carry) = self.state[12].overflowing_add(1);
173 self.state[12] = c0;
174 if self.wide_counter && carry {
175 self.state[13] = self.state[13].wrapping_add(1);
176 }
177 }
178}
179
180impl Cipher for ChaCha20 {
181 fn process(&mut self, data: &mut [u8]) {
182 for b in data.iter_mut() {
183 if self.pos == 64 {
184 self.refill();
185 }
186 *b ^= self.block[self.pos];
187 self.pos += 1;
188 }
189 }
190}
191
192fn chacha_key(session_key: &[u8]) -> [u8; 32] {
194 let mut h = Sha256::new();
195 h.update(session_key);
196 h.finalize().into()
197}
198
199pub fn make_ciphers(
205 plugin: WireCryptPlugin,
206 key: &[u8],
207 nonce: &[u8],
208) -> (Box<dyn Cipher>, Box<dyn Cipher>) {
209 match plugin {
210 WireCryptPlugin::Arc4 => (Box::new(Rc4::new(key)), Box::new(Rc4::new(key))),
211 WireCryptPlugin::ChaCha | WireCryptPlugin::ChaCha64 => {
212 let k = chacha_key(key);
213 (
214 Box::new(ChaCha20::new(&k, nonce)),
215 Box::new(ChaCha20::new(&k, nonce)),
216 )
217 }
218 }
219}
220
221#[cfg(test)]
222mod tests {
223 use super::*;
224
225 #[test]
227 fn rc4_known_answer() {
228 let mut c = Rc4::new(b"Key");
230 let mut data = b"Plaintext".to_vec();
231 c.process(&mut data);
232 assert_eq!(data, hex_to_vec("BBF316E8D940AF0AD3"));
233 }
234
235 #[test]
236 fn rc4_roundtrip() {
237 let key = b"firebird-session-key";
238 let mut enc = Rc4::new(key);
239 let mut dec = Rc4::new(key);
240 let mut buf = b"op_attach payload \x00\x01\x02".to_vec();
241 let orig = buf.clone();
242 enc.process(&mut buf);
243 assert_ne!(buf, orig);
244 dec.process(&mut buf);
245 assert_eq!(buf, orig);
246 }
247
248 fn hex_to_vec(s: &str) -> Vec<u8> {
249 let s: String = s.chars().filter(|c| !c.is_whitespace()).collect();
250 (0..s.len())
251 .step_by(2)
252 .map(|i| u8::from_str_radix(&s[i..i + 2], 16).unwrap())
253 .collect()
254 }
255
256 #[test]
261 fn chacha20_rfc8439_block() {
262 let key: Vec<u8> = (0u8..32).collect();
263 let nonce = hex_to_vec("000000090000004a00000000");
264 let expected = hex_to_vec(
265 "10f1e7e4d13b5915500fdd1fa32071c4
266 c7d1f4c733c0680304 22aa9ac3d46c4e
267 d2826446079faa0914c2d705d98b02a2
268 b5129cd1de164eb9cbd083e8a2503c4e",
269 );
270 let mut c = ChaCha20::new(&key, &nonce);
271 let mut buf = vec![0u8; 128];
272 c.process(&mut buf);
273 assert_eq!(&buf[64..128], &expected[..]);
274 }
275
276 #[test]
277 fn chacha20_roundtrip_both_variants() {
278 let key = [0x42u8; 32];
279 for nonce in [vec![7u8; 12], vec![9u8; 8]] {
280 let mut enc = ChaCha20::new(&key, &nonce);
281 let mut dec = ChaCha20::new(&key, &nonce);
282 let orig = b"op_attach + segredo \x00\x01\x02 atravessa varios blocos ".repeat(4);
283 let mut buf = orig.clone();
284 enc.process(&mut buf);
285 assert_ne!(buf, orig);
286 dec.process(&mut buf);
287 assert_eq!(buf, orig);
288 }
289 }
290}