1use crate::{Asset, Error, Result};
6use flate2::write::GzEncoder;
7use flate2::Compression;
8use std::collections::HashSet;
9use std::fs::File;
10use std::io::Write;
11use std::path::Path;
12use tar::{Builder, Header};
13
14#[derive(Debug, Default)]
16pub struct UnityPackageBuilder {
17 assets: Vec<Asset>,
18 deterministic_guids: bool,
19}
20
21impl UnityPackageBuilder {
22 pub fn new() -> Self {
24 Self {
25 assets: Vec::new(),
26 deterministic_guids: true,
27 }
28 }
29
30 pub fn asset(mut self, asset: Asset) -> Self {
32 self.assets.push(asset);
33 self
34 }
35
36 pub fn assets(mut self, assets: impl IntoIterator<Item = Asset>) -> Self {
38 self.assets.extend(assets);
39 self
40 }
41
42 pub fn deterministic_guids(mut self, enabled: bool) -> Self {
44 self.deterministic_guids = enabled;
45 self
46 }
47
48 pub fn build(self) -> Result<UnityPackage> {
50 let mut pkg = UnityPackage::new();
51
52 for mut asset in self.assets {
53 if !self.deterministic_guids {
54 asset = asset.with_random_guid();
55 }
56 pkg.add(asset)?;
57 }
58
59 Ok(pkg)
60 }
61}
62
63#[derive(Debug, Default)]
65pub struct UnityPackage {
66 assets: Vec<Asset>,
67 paths: HashSet<String>,
68}
69
70impl UnityPackage {
71 pub fn new() -> Self {
73 Self::default()
74 }
75
76 pub fn builder() -> UnityPackageBuilder {
78 UnityPackageBuilder::new()
79 }
80
81 pub fn add(&mut self, asset: Asset) -> Result<&mut Self> {
83 if !asset.path.starts_with("Assets/") && !asset.path.starts_with("Assets\\") {
85 return Err(Error::PathMustStartWithAssets(asset.path.clone()));
86 }
87
88 if self.paths.contains(&asset.path) {
90 return Err(Error::DuplicatePath(asset.path.clone()));
91 }
92
93 self.paths.insert(asset.path.clone());
94 self.assets.push(asset);
95 Ok(self)
96 }
97
98 pub fn add_binary(&mut self, path: impl Into<String>, data: Vec<u8>) -> Result<&mut Self> {
100 self.add(Asset::binary(path, data))
101 }
102
103 pub fn add_text(
105 &mut self,
106 path: impl Into<String>,
107 content: impl Into<String>,
108 ) -> Result<&mut Self> {
109 self.add(Asset::text(path, content))
110 }
111
112 pub fn add_folder(&mut self, path: impl Into<String>) -> Result<&mut Self> {
114 self.add(Asset::folder(path))
115 }
116
117 pub fn assets(&self) -> &[Asset] {
119 &self.assets
120 }
121
122 pub fn len(&self) -> usize {
124 self.assets.len()
125 }
126
127 pub fn is_empty(&self) -> bool {
129 self.assets.is_empty()
130 }
131
132 pub fn save(&self, path: impl AsRef<Path>) -> Result<()> {
134 let file = File::create(path)?;
135 self.write_to(file)
136 }
137
138 pub fn write_to<W: Write>(&self, writer: W) -> Result<()> {
140 let encoder = GzEncoder::new(writer, Compression::default());
142 let mut tar = Builder::new(encoder);
143
144 let mut folders_to_create: HashSet<String> = HashSet::new();
146 for asset in &self.assets {
147 let path = &asset.path;
148 let mut current = String::new();
149 for component in path.split('/').take(path.split('/').count() - 1) {
150 if current.is_empty() {
151 current = component.to_string();
152 } else {
153 current = format!("{}/{}", current, component);
154 }
155 if current.starts_with("Assets") && !self.paths.contains(¤t) {
156 folders_to_create.insert(current.clone());
157 }
158 }
159 }
160
161 let folder_assets: Vec<Asset> = folders_to_create.into_iter().map(Asset::folder).collect();
163
164 let all_assets: Vec<&Asset> = folder_assets.iter().chain(self.assets.iter()).collect();
166
167 for asset in all_assets {
169 self.write_asset(&mut tar, asset)?;
170 }
171
172 let encoder = tar.into_inner()?;
174 encoder.finish()?;
175
176 Ok(())
177 }
178
179 fn write_asset<W: Write>(&self, tar: &mut Builder<W>, asset: &Asset) -> Result<()> {
181 let guid_str = asset.guid.to_hex();
182
183 let pathname_path = format!("{}/pathname", guid_str);
185 let pathname_content = format!("{}\n", asset.path);
186 self.write_tar_file(tar, &pathname_path, pathname_content.as_bytes())?;
187
188 if let Some(content) = asset.content_bytes() {
190 let asset_path = format!("{}/asset", guid_str);
191 self.write_tar_file(tar, &asset_path, &content)?;
192 }
193
194 let meta_path = format!("{}/asset.meta", guid_str);
196 let meta_content = asset.meta_file().to_yaml();
197 self.write_tar_file(tar, &meta_path, meta_content.as_bytes())?;
198
199 Ok(())
200 }
201
202 fn write_tar_file<W: Write>(
204 &self,
205 tar: &mut Builder<W>,
206 path: &str,
207 data: &[u8],
208 ) -> Result<()> {
209 let mut header = Header::new_gnu();
210 header.set_size(data.len() as u64);
211 header.set_mode(0o644);
212 header.set_mtime(0); header.set_cksum();
214
215 tar.append_data(&mut header, path, data)?;
216 Ok(())
217 }
218}
219
220#[cfg(test)]
221mod tests {
222 use super::*;
223
224 #[test]
225 fn test_empty_package() {
226 let pkg = UnityPackage::new();
227 assert!(pkg.is_empty());
228 assert_eq!(pkg.len(), 0);
229 }
230
231 #[test]
232 fn test_add_assets() {
233 let mut pkg = UnityPackage::new();
234 pkg.add_binary("Assets/test.raw", vec![1, 2, 3]).unwrap();
235 pkg.add_text("Assets/config.json", "{}").unwrap();
236
237 assert_eq!(pkg.len(), 2);
238 }
239
240 #[test]
241 fn test_path_validation() {
242 let mut pkg = UnityPackage::new();
243
244 assert!(pkg.add_binary("Assets/test.raw", vec![]).is_ok());
246
247 let result = pkg.add_binary("test.raw", vec![]);
249 assert!(matches!(result, Err(Error::PathMustStartWithAssets(_))));
250 }
251
252 #[test]
253 fn test_duplicate_path() {
254 let mut pkg = UnityPackage::new();
255 pkg.add_binary("Assets/test.raw", vec![1]).unwrap();
256
257 let result = pkg.add_binary("Assets/test.raw", vec![2]);
258 assert!(matches!(result, Err(Error::DuplicatePath(_))));
259 }
260
261 #[test]
262 fn test_write_package() {
263 let mut pkg = UnityPackage::new();
264 pkg.add_binary("Assets/Terrain/heightmap.raw", vec![0, 1, 2, 3])
265 .unwrap();
266 pkg.add_text("Assets/Terrain/metadata.json", r#"{"width": 256}"#)
267 .unwrap();
268
269 let mut buffer = Vec::new();
270 pkg.write_to(&mut buffer).unwrap();
271
272 assert!(buffer.len() > 0);
274 assert_eq!(&buffer[0..2], &[0x1f, 0x8b]); }
276
277 #[test]
278 fn test_builder_basic() {
279 let pkg = UnityPackage::builder()
280 .asset(Asset::binary("Assets/test.raw", vec![1, 2, 3]))
281 .asset(Asset::text("Assets/config.json", "{}"))
282 .build()
283 .unwrap();
284
285 assert_eq!(pkg.len(), 2);
286 }
287
288 #[test]
289 fn test_builder_deterministic_guids() {
290 let pkg1 = UnityPackage::builder()
291 .asset(Asset::binary("Assets/test.raw", vec![1]))
292 .build()
293 .unwrap();
294
295 let pkg2 = UnityPackage::builder()
296 .asset(Asset::binary("Assets/test.raw", vec![2]))
297 .build()
298 .unwrap();
299
300 assert_eq!(pkg1.assets()[0].guid, pkg2.assets()[0].guid);
302 }
303
304 #[test]
305 fn test_builder_random_guids() {
306 let pkg1 = UnityPackage::builder()
307 .asset(Asset::binary("Assets/test.raw", vec![1]))
308 .deterministic_guids(false)
309 .build()
310 .unwrap();
311
312 let pkg2 = UnityPackage::builder()
313 .asset(Asset::binary("Assets/test.raw", vec![2]))
314 .deterministic_guids(false)
315 .build()
316 .unwrap();
317
318 assert_ne!(pkg1.assets()[0].guid, pkg2.assets()[0].guid);
320 }
321}