lolistealer/src/loli/mod.rs

use std::{
    path::Path,
    convert::TryInto,
    fs::{
	File,
	OpenOptions,
    },
    ops::{
	Range,
    },
};
use memmap::{
    MmapOptions,
    Mmap,
    MmapMut,
};
pub mod error;
pub mod encoding;

/// Attempt to decode an image
fn decode<S, W>(input: S, mut output: W) -> Result<usize, error::Error>
where S: AsRef<[u8]>,
      W: AsMut<[u8]>
{
    let input_bytes = input.as_ref();
    let output_bytes = output.as_mut();

    //panic!("{:?}", std::str::from_utf8(input_bytes).expect("XXX fatal"));
    Ok(base64::decode_config_slice(input_bytes, base64::STANDARD, output_bytes)?) // XXX: This fails
}

/// Calculate the size for a base64 inpue
fn calc_size<T>(input: T) -> usize
where T: AsRef<[u8]>
{
    encoding::data_size(input.as_ref().len())
}

/// Try to get the image format from a `str` slice.
#[inline]
fn attempt_get_format<T>(input: T) -> Result<encoding::ImageType, error::Error>
where T: AsRef<str>
{
    match input.as_ref().parse() {
	Ok(v) => Ok(v),
	Err(error::Error::UnknownFormat) => Ok(Default::default()),
	Err(e) => Err(e),
    }
}

/// An encoded loli image
#[derive(Debug)]
pub struct BasedLoli
{
    size_decoded: usize,
    map: Mmap, // We should probably have this dropped before `file`
    file: File,
}

impl BasedLoli
{
    /// Create a new map to basedloli.
    pub fn map(file: impl AsRef<Path>) -> Result<Self, error::Error>
    {
	let file = File::open(file.as_ref())?;

	let meta = file.metadata()?;
	Ok(Self {
	    size_decoded: meta.len().try_into()?,
	    map: unsafe { MmapOptions::new().map(&file)? },
	    file,
	})
    }

    /// Find bounds for base64 data
    pub fn calculate_bounds<'a>(&'a self) -> Result<LoliBounds<'a>, error::Error>
    {
	let bound = encoding::find_bounds(self.as_ref())?;
	let image_type = self.try_get_type(bound.start)?;
	let tags_range = encoding::find_tag_bounds(&self.map[bound.end..]).ok();
	Ok(LoliBounds {
	    loli: self,
	    image_type,
	    range:bound,
	    tags_range,
	})
    }

    /// Get the raw bytes of this map
    pub fn bytes(&self) -> &[u8]
    {
	self.map.as_ref()
    }

    /// Try to get as a str
    pub fn try_as_str(&self) -> Result<&str, error::DecodeError>
    {
	Ok(std::str::from_utf8(self.map.as_ref())?)
    }

    /// Try to get the image type from encoded data
    pub fn try_get_type(&self, start: usize) -> Result<encoding::ImageType, error::Error>
    {
	attempt_get_format(&(self.try_as_str()?)[start..])
    }

    /// The calculated size
    pub fn decoded_size(&self) -> usize
    {
	self.size_decoded
    }

}

impl<'a> LoliBounds<'a>
{

    /// Get the type of image
    pub fn image(&self) -> &encoding::ImageType
    {
	&self.image_type
    }

    pub fn bounds(&self) -> &Range<usize>
    {
	&self.range
    }
    
    /// Try to create a decoding container for this `BasedLoli` to the file in `path`.
    pub fn create_child(&self, path: impl AsRef<Path>) -> Result<Loli, error::Error>
    {
	let size = self.loli.size_decoded;
	let image_type = self.image_type.clone();
	let path = path.as_ref();
	
	let file = OpenOptions::new()
	    .create(true)
	    .write(true)
	    .read(true)
	    .open(path)?;
	file.set_len(size.try_into()?)?;

	Ok(Loli {
	    size,
	    image_type,
	    map: unsafe { MmapOptions::new().map_mut(&file).map_err(|e| error::DecodeError::Map(e, path.into()))? },
	    file: file,
	    tags: if let Some(range) = &self.tags_range {
		let bytes = self.loli.bytes();
		let range = &bytes[(self.range.end+range.start)..(self.range.end+range.end)];
		if let Ok(string) = std::str::from_utf8(range) {
		    string.split(' ').map(|x| x.to_owned()).collect()
		} else {
		    Vec::default()
		}
	    } else { Vec::default() },
	})
    }    
    /// Attempt to decode to a child container of our owner
    #[inline]
    pub fn decode(&self, loli: &mut Loli) -> Result<usize, error::Error>
    {
	let bytes = self.loli.bytes();

	decode(&bytes[self.range.clone()], loli)
    }
}

impl AsRef<[u8]> for BasedLoli
{
    fn as_ref(&self) -> &[u8]
    {
	self.bytes()
    }
}

#[derive(Debug)]
pub struct Loli
{
    size: usize,
    image_type: encoding::ImageType,
    map: MmapMut,
    file: File,
    tags: Vec<String>,
}

impl AsMut<[u8]> for Loli
{
    fn as_mut(&mut self) -> &mut [u8]
    {
	self.map.as_mut()
    }
}

#[derive(Debug)]
pub struct LoliBounds<'a>
{
    loli: &'a BasedLoli,
    range: Range<usize>,
    tags_range: Option<Range<usize>>,
    image_type: encoding::ImageType,
}

impl Loli {
    /// Get the tags for this loli
    pub fn tags(&self) -> &[String] //TODO: Tags filter
    {
	&self.tags[..]
    }
}