1  use std::io;
  2  
  3  use deku::{prelude::*, reader::ReaderRet};
  4  use log::*;
  5  use serde::{Deserialize, Serialize};
  6  
  7  #[cfg(feature = "python")]
  8  use pyo3::prelude::*;
  9  
 10  use crate::VecReader;
 11  
 12  pub const MAGIC: u32 = 0xAFBF0C01;
 13  
 14  #[derive(Debug, Clone, Default)]
 15  pub struct Context {
 16  	/// Whether to skip the `Rpk` data field or not.
 17  	///
 18  	/// By default, the data field will be read by the reader. If you want to save memory and have
 19  	/// better performance, and the data field is not needed such as when simply viewing a list of
 20  	/// entries, setting this to `true` may be desired.
 21  	pub entries_only: bool,
 22  	/// Filter the `Rpk` data field by a desired `Vec<Entry>`.
 23  	///
 24  	/// By default, no filter is applied thus all of the data is read by the reader. If only data
 25  	/// from certain entries is desired, provide this field with your desired entries.
 26  	pub entries: Option<Vec<Entry>>,
 27  }
 28  
 29  #[cfg(feature = "python")]
 30  /// Rayform Package
 31  #[pyclass]
 32  #[deku_derive(DekuRead, DekuWrite)]
 33  #[derive(Clone, Debug, Deserialize, Serialize)]
 34  #[deku(ctx = "ctx: Context", ctx_default = "Context::default()")]
 35  pub struct Rpk {
 36  	#[deku(temp, temp_value = "(entries.len() * 32) as u32")]
 37  	pub size_of_entries: u32,
 38  	#[pyo3(get)]
 39  	#[deku(count = "size_of_entries / 32")]
 40  	pub entries: Vec<Entry>,
 41  	#[pyo3(get)]
 42  	#[deku(
 43  		cond = "!ctx.entries_only",
 44  		reader = "Rpk::read(ctx, deku::reader, entries)",
 45  		writer = "VecReader::write_nested(deku::writer, &self.data)"
 46  	)]
 47  	pub data: Vec<Vec<u8>>,
 48  }
 49  
 50  #[cfg(not(feature = "python"))]
 51  /// Rayform Package
 52  #[deku_derive(DekuRead, DekuWrite)]
 53  #[derive(Clone, Debug, Deserialize, Serialize)]
 54  #[deku(ctx = "ctx: Context", ctx_default = "Context::default()")]
 55  pub struct Rpk {
 56  	#[deku(temp, temp_value = "(entries.len() * 32) as u32")]
 57  	pub size_of_entries: u32,
 58  	#[deku(count = "size_of_entries / 32")]
 59  	pub entries: Vec<Entry>,
 60  	#[deku(
 61  		cond = "!ctx.entries_only",
 62  		reader = "Rpk::read(ctx, deku::reader, entries)",
 63  		writer = "VecReader::write_nested(deku::writer, &self.data)"
 64  	)]
 65  	pub data: Vec<Vec<u8>>,
 66  }
 67  
 68  impl Rpk {
 69  	fn read<R: io::Read + io::Seek>(
 70  		ctx: Context,
 71  		reader: &mut Reader<R>,
 72  		entries: &[Entry],
 73  	) -> Result<Vec<Vec<u8>>, DekuError> {
 74  		// Sort the entries by offset so we can read them in order.
 75  		let mut entries = match ctx.entries {
 76  			Some(entries) => entries,
 77  			None => entries.to_vec(),
 78  		};
 79  		entries.sort_by(|a, b| a.offset.cmp(&b.offset));
 80  
 81  		let mut formats: Vec<Vec<u8>> = Vec::with_capacity(entries.len());
 82  
 83  		for entry in entries {
 84  			let mut buf = vec![0; entry.size as usize];
 85  
 86  			// deku is slow at reading Vec<u8> so we read the bytes into a buffer ourselves.
 87  			let format = match reader.read_bytes(entry.size as usize, &mut buf) {
 88  				Ok(ReaderRet::Bytes) => Ok(buf),
 89  				Ok(ReaderRet::Bits(_)) => {
 90  					Err(DekuError::InvalidParam("Expected bytes, got bits".into()))
 91  				}
 92  				Err(err) => Err(err),
 93  			}?;
 94  
 95  			formats.push(format);
 96  		}
 97  		Ok(formats)
 98  	}
 99  }
100  
101  #[cfg_attr(feature = "python", pyclass(get_all))]
102  #[derive(Debug, DekuRead, DekuWrite, Deserialize, Serialize, Clone)]
103  pub struct Entry {
104  	#[deku(
105  		reader = "Entry::read(deku::reader)",
106  		writer = "Entry::write(deku::writer, &self.name)"
107  	)]
108  	pub name: String,
109  	pub offset: u32,
110  	#[deku(pad_bytes_after = "8")]
111  	pub size: u32,
112  }
113  
114  impl Entry {
115  	fn read<R: io::Read + io::Seek>(reader: &mut Reader<R>) -> Result<String, DekuError> {
116  		let mut buf: Vec<u8> = vec![0; 16];
117  		match reader.read_bytes(16, &mut buf) {
118  			Ok(ReaderRet::Bytes) => {
119  				let mut s = String::new();
120  				for b in buf {
121  					match b {
122  						0 => break,
123  						v => s.push(v as char),
124  					}
125  				}
126  				Ok(s)
127  			}
128  			Ok(ReaderRet::Bits(_)) => {
129  				Err(DekuError::InvalidParam("Expected bytes, got bits".into()))
130  			}
131  			Err(err) => Err(err),
132  		}
133  	}
134  
135  	fn write<W: io::Write + io::Seek>(
136  		writer: &mut Writer<W>,
137  		value: &str,
138  	) -> Result<(), DekuError> {
139  		let mut buf = vec![0; 16];
140  		for (i, c) in value.chars().enumerate() {
141  			if i >= 16 {
142  				warn!(
143  					r#"entry name "{}" has been truncated due to being longer than 16 characters."#,
144  					value
145  				);
146  				break;
147  			}
148  			buf[i] = c as u8;
149  		}
150  		writer.write_bytes(&buf)
151  	}
152  }