libtlg-rs 0.2.4

Rust version of libtlg
Documentation
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()
    }
}