use std::collections::HashMap;
#[cfg(feature = "encode")]
use std::io::Write;
use std::io::{Read, Seek, SeekFrom};
pub trait ReadExt {
fn read_u32(&mut self) -> std::io::Result<u32>;
fn read_u8(&mut self) -> std::io::Result<u8>;
}
impl<R: Read> ReadExt for R {
fn read_u32(&mut self) -> std::io::Result<u32> {
let mut buf = [0; 4];
self.read_exact(&mut buf)?;
Ok(u32::from_le_bytes(buf))
}
fn read_u8(&mut self) -> std::io::Result<u8> {
let mut buf = [0; 1];
self.read_exact(&mut buf)?;
Ok(buf[0])
}
}
#[cfg(feature = "encode")]
pub trait WriteExt {
fn write_u32(&mut self, value: u32) -> std::io::Result<()>;
fn write_u8(&mut self, value: u8) -> std::io::Result<()>;
}
#[cfg(feature = "encode")]
impl<W: Write> WriteExt for W {
fn write_u32(&mut self, value: u32) -> std::io::Result<()> {
self.write_all(&value.to_le_bytes())
}
fn write_u8(&mut self, value: u8) -> std::io::Result<()> {
self.write_all(&[value])
}
}
pub struct XorStream<R: Read + Seek> {
inner: R,
keys: HashMap<u64, u8>,
}
impl<R: Read + Seek> XorStream<R> {
pub fn new(inner: R) -> Self {
Self {
inner,
keys: HashMap::new(),
}
}
pub fn set_key(&mut self, offset: u64, key: u8) {
self.keys.insert(offset, key);
}
}
impl<R: Read + Seek> Read for XorStream<R> {
fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
let pos = self.inner.stream_position()?;
let bytes_read = self.inner.read(buf)?;
for i in 0..bytes_read {
if let Some(key) = self.keys.get(&(pos + i as u64)) {
buf[i] ^= *key;
}
}
Ok(bytes_read)
}
}
impl<R: Read + Seek> Seek for XorStream<R> {
fn seek(&mut self, pos: SeekFrom) -> std::io::Result<u64> {
self.inner.seek(pos)
}
fn rewind(&mut self) -> std::io::Result<()> {
self.inner.rewind()
}
fn seek_relative(&mut self, offset: i64) -> std::io::Result<()> {
self.inner.seek_relative(offset)
}
fn stream_position(&mut self) -> std::io::Result<u64> {
self.inner.stream_position()
}
}