importer.rs (21894B)
1 use std::path::PathBuf; 2 use std::time::Duration; 3 4 use async_nats::jetstream; 5 use async_nats::jetstream::consumer::pull::Config; 6 use async_nats::jetstream::consumer::Consumer; 7 use async_nats::jetstream::{AckKind, Context}; 8 use futures::StreamExt; 9 use notify_debouncer_full::notify::RecommendedWatcher; 10 use notify_debouncer_full::{new_debouncer, DebounceEventResult, Debouncer, FileIdMap}; 11 use rexiv2::Metadata; 12 use time::macros::format_description; 13 use time::PrimitiveDateTime; 14 use tokio::runtime::Handle; 15 use tracing::{debug, error, info, trace, warn}; 16 17 use crate::jetstream_events::Exif; 18 19 const ALLOWED_TYPES: [&str; 8] = ["cr2", "cr3", "jpg", "tiff", "dng", "png", "jpeg", "mp4"]; 20 21 #[derive(Debug, Clone)] 22 pub(crate) struct Importer { 23 // Folder to watch for new files 24 import_folder: PathBuf, 25 // Export folder to move files to 26 export_folder: PathBuf, 27 // Nats consumer to work with 28 consumer: Consumer<Config>, 29 // Jetstream client 30 jetstream: Context, 31 } 32 33 impl Importer { 34 pub(crate) async fn new( 35 import_folder: PathBuf, 36 export_folder: PathBuf, 37 import_stream: jetstream::stream::Stream, 38 jetstream: Context, 39 ) -> color_eyre::Result<Self> { 40 let consumer = import_stream 41 .create_consumer(Config { 42 durable_name: Some("processor-1".to_string()), 43 filter_subject: "import.file_added".to_string(), 44 ..Default::default() 45 }) 46 .await?; 47 48 Ok(Self { 49 import_folder, 50 export_folder, 51 consumer, 52 jetstream, 53 }) 54 } 55 56 fn extract_exif<'a>(&self, exif: Metadata) -> Exif<'a> { 57 // Get all the exifdata 58 let gps_info = exif.get_gps_info(); 59 let gps_latitude_ref = exif.get_tag_string("Exif.GPSInfo.GPSLatitudeRef"); 60 let gps_longitude_ref = exif.get_tag_string("Exif.GPSInfo.GPSLongitudeRef"); 61 let gps_timestamp = exif.get_tag_rational("Exif.GPSInfo.GPSTimeStamp"); 62 let gps_speed_ref = exif.get_tag_string("Exif.GPSInfo.GPSSpeedRef"); 63 let gps_speed = exif.get_tag_rational("Exif.GPSInfo.GPSSpeed"); 64 let gps_direction_ref = exif.get_tag_string("Exif.GPSInfo.GPSImgDirectionRef"); 65 let gps_direction = exif.get_tag_rational("Exif.GPSInfo.GPSImgDirection"); 66 let gps_positioning_error = exif.get_tag_rational("Exif.GPSInfo.GPSHPositioningError"); 67 let exposure_time = exif.get_tag_rational("Exif.Photo.ExposureTime"); 68 let f_number = exif.get_tag_rational("Exif.Photo.FNumber"); 69 let digital_zoom_ratio = exif.get_tag_rational("Exif.Photo.DigitalZoomRatio"); 70 let lens_make = exif.get_tag_string("Exif.Photo.LensMake"); 71 let lens_model = exif.get_tag_string("Exif.Photo.LensModel"); 72 73 Exif { 74 gps_version_id: None, 75 latitude: gps_info.map(|info| { 76 format!( 77 "{}{}", 78 info.latitude.to_string().replace('.', ","), 79 gps_latitude_ref.unwrap_or_default() 80 ) 81 .into() 82 }), 83 longitude: gps_info.map(|info| { 84 format!( 85 "{}{}", 86 info.longitude.to_string().replace('.', ","), 87 gps_longitude_ref.unwrap_or_default() 88 ) 89 .into() 90 }), 91 altitude_ref: None, 92 altitude: gps_info.map(|info| info.altitude.to_string().into()), 93 timestamp: gps_timestamp.map(|timestamp| timestamp.to_string().into()), 94 speed_ref: gps_speed_ref.map(|speed_ref| speed_ref.into()).ok(), 95 speed: gps_speed.map(|speed| speed.to_string().into()), 96 direction_ref: gps_direction_ref 97 .map(|direction_ref| direction_ref.into()) 98 .ok(), 99 direction: gps_direction.map(|direction| direction.to_string().into()), 100 destination_bearing_ref: None, 101 destination_bearing: None, 102 positioning_error: gps_positioning_error.map(|error| error.to_string().into()), 103 exposure_time: exposure_time.map(|time| time.to_string().into()), 104 // Cursed 105 f_number: f_number 106 .map(|f_number| f_number.into_raw().0 as f64 / f_number.into_raw().1 as f64), 107 color_space: None, 108 digital_zoom_ratio: digital_zoom_ratio.map(|digital_zoom_ratio| { 109 digital_zoom_ratio.into_raw().0 as f64 / digital_zoom_ratio.into_raw().1 as f64 110 }), 111 make: None, 112 model: None, 113 lens_make: lens_make.map(|lens_make| lens_make.into()).ok(), 114 lens_model: lens_model.map(|lens_model| lens_model.into()).ok(), 115 lens_specification: None, 116 } 117 } 118 119 pub(crate) async fn nats_consumer(&self) -> color_eyre::Result<()> { 120 let mut messages = self.consumer.messages().await?; 121 while let Some(msg) = messages.next().await { 122 match msg { 123 Ok(msg) => { 124 let event: super::jetstream_events::FileToImport = 125 bincode::deserialize(&msg.payload).unwrap(); 126 debug!("Received message: {:#?}", event); 127 128 // Initial ack to notify that we received the message 129 msg.ack_with(AckKind::Progress).await.unwrap(); 130 131 // Check if the file still is in the import 132 133 // First get absolute import folder path 134 let import_folder = self.import_folder.canonicalize()?; 135 136 // Check if the file is still there 137 if !event.path.exists() { 138 error!("[IMPORTER] file does not exist: {:?}", event.path); 139 // We also acknowledge the message to remove it from the queue 140 msg.ack_with(AckKind::Term).await.unwrap(); 141 continue; 142 } 143 144 // Then check if the file is in the import folder 145 if !event.path.starts_with(import_folder.clone()) { 146 error!( 147 "[IMPORTER] file is not in the import folder: {:?}", 148 event.path 149 ); 150 // Move the folder into the "failed" subfolder of the import folder 151 let failed_folder = import_folder.join("failed"); 152 std::fs::create_dir_all(&failed_folder)?; 153 let new_path = failed_folder.join(event.path.file_name().unwrap()); 154 std::fs::rename(&event.path, &new_path)?; 155 156 // Also move any xmp files with the same name 157 let xmp_path = event.path.with_extension("xmp"); 158 if xmp_path.exists() { 159 let new_xmp_path = failed_folder.join(xmp_path.file_name().unwrap()); 160 std::fs::rename(&xmp_path, &new_xmp_path)?; 161 } 162 163 // We also acknowledge the message to remove it from the queue 164 msg.ack_with(AckKind::Term).await.unwrap(); 165 continue; 166 } 167 168 // Another progress ack to notify that we are processing the file 169 msg.ack_with(AckKind::Progress).await.unwrap(); 170 171 let mut result_path = event.path.clone(); 172 173 let mut exif_data = None; 174 // Get exif data 175 { 176 let meta = rexiv2::Metadata::new_from_path(&event.path); 177 178 match meta { 179 Ok(exif) => { 180 trace!("Got exif data"); 181 // Get the date the photo was taken 182 if let Ok(date) = exif.get_tag_string("Exif.Photo.DateTimeOriginal") 183 { 184 let format = format_description!( 185 "[year]:[month]:[day] [hour]:[minute]:[second]" 186 ); 187 match PrimitiveDateTime::parse(&date, &format) { 188 Ok(datetime) => { 189 debug!("[IMPORTER] Date taken: {:?}", datetime); 190 // Ensure the export subfolders exist 191 // The image should be saved in `export/year/month/day` format 192 let year = datetime.year().to_string(); 193 let month = datetime.month().to_string(); 194 let day = datetime.day().to_string(); 195 let export_folder = self 196 .export_folder 197 .join(&year) 198 .join(&month) 199 .join(&day); 200 std::fs::create_dir_all(&export_folder)?; 201 202 // Move the file to the export folder 203 let new_path = 204 export_folder.join(event.path.file_name().unwrap()); 205 std::fs::rename(&event.path, &new_path)?; 206 207 // Also move any xmp files with the same name 208 let xmp_path = event.path.with_extension("xmp"); 209 if xmp_path.exists() { 210 let new_xmp_path = export_folder 211 .join(xmp_path.file_name().unwrap()); 212 std::fs::rename(&xmp_path, &new_xmp_path)?; 213 } 214 215 result_path = new_path; 216 } 217 Err(e) => { 218 error!("[IMPORTER] error parsing date: {:?}", e); 219 } 220 } 221 } else { 222 let other_folder = self.export_folder.join("other"); 223 error!( 224 "[IMPORTER] Date taken not found placing in \"{}\"", 225 other_folder.display() 226 ); 227 // Ensure the export "other" subfolder exists 228 std::fs::create_dir_all(&other_folder)?; 229 230 // Move the file to the export folder 231 let new_path = 232 other_folder.join(event.path.file_name().unwrap()); 233 std::fs::rename(&event.path, &new_path)?; 234 235 // Also move any xmp files with the same name 236 let xmp_path = event.path.with_extension("xmp"); 237 if xmp_path.exists() { 238 let new_xmp_path = 239 other_folder.join(xmp_path.file_name().unwrap()); 240 std::fs::rename(&xmp_path, &new_xmp_path)?; 241 } 242 243 result_path = new_path; 244 } 245 246 let exif = self.extract_exif(exif); 247 248 exif_data = Some(exif); 249 250 info!("Finished processing exif data") 251 } 252 Err(e) => { 253 error!("[IMPORTER] error reading exif data: {:?}", e); 254 let other_folder = self.export_folder.join("other"); 255 error!( 256 "[IMPORTER] Exif data not found placing in \"{}\"", 257 other_folder.display() 258 ); 259 // Ensure the export "other" subfolder exists 260 std::fs::create_dir_all(&other_folder)?; 261 262 // Move the file to the export folder 263 let new_path = other_folder.join(event.path.file_name().unwrap()); 264 std::fs::rename(&event.path, &new_path)?; 265 266 // Also move any xmp files with the same name 267 let xmp_path = event.path.with_extension("xmp"); 268 if xmp_path.exists() { 269 let new_xmp_path = 270 other_folder.join(xmp_path.file_name().unwrap()); 271 std::fs::rename(&xmp_path, &new_xmp_path)?; 272 } 273 274 result_path = new_path; 275 } 276 } 277 }; 278 279 info!( 280 "Finished processing file: {:?}\nQueuing for next steps", 281 result_path 282 ); 283 284 // Add another progress ack to notify that we are done processing the exif 285 msg.ack_with(AckKind::Progress).await.unwrap(); 286 287 // Add a c2pa event to the stream 288 let c2pa_event = super::jetstream_events::C2PASignRequest { 289 path: result_path.canonicalize().unwrap(), 290 exif: exif_data, 291 }; 292 let bytes = bincode::serialize(&c2pa_event).unwrap(); 293 publish_with_error_printing( 294 &msg.context, 295 "import.sign_request".to_string(), 296 bytes, 297 ) 298 .await; 299 300 // Add another progress ack to notify that we are done processing the exif 301 msg.ack_with(AckKind::Progress).await.unwrap(); 302 303 // TODO: Notify other steps here? 304 305 // Acknowledge the message to remove it from the queue 306 msg.ack_with(AckKind::Ack).await.unwrap(); 307 tokio::time::sleep(std::time::Duration::from_millis(100)).await; 308 } 309 Err(e) => { 310 error!("[IMPORTER] error receiving message: {:?}", e); 311 } 312 } 313 } 314 error!("[IMPORTER] Consumer stopped"); 315 Ok(()) 316 } 317 318 async fn file_added(path: PathBuf, jetstream: Context) -> color_eyre::Result<()> { 319 // Check if the path is a file 320 if !path.is_file() { 321 return Ok(()); 322 } 323 // check if cr2, cr3, jpg, tiff, dng or png file 324 let extension = path 325 .extension() 326 .unwrap_or_default() 327 .to_string_lossy() 328 .to_lowercase(); 329 if !ALLOWED_TYPES.contains(&extension.as_str()) { 330 return Ok(()); 331 } 332 info!("[IMPORTER] New file to import: {:?}", path); 333 334 // Notify the importer to import the file 335 let import_event = super::jetstream_events::FileToImport { 336 path: path.canonicalize().unwrap(), 337 }; 338 let bytes = bincode::serialize(&import_event).unwrap(); 339 340 publish_with_error_printing(&jetstream, "import.file_added".to_string(), bytes).await; 341 342 Ok(()) 343 } 344 345 fn file_event_handler(res: DebounceEventResult, jetstream: Context, tokio_runtime: Handle) { 346 match res { 347 Ok(events) => { 348 for event in events { 349 // Check if the file was copied to and not removed 350 if event.kind.is_remove() { 351 continue; 352 } 353 for path in event.paths.clone() { 354 let jetstream = jetstream.clone(); 355 tokio_runtime.spawn(async move { 356 Importer::file_added(path, jetstream).await.unwrap(); 357 tokio::time::sleep(Duration::from_millis(100)).await; 358 }); 359 } 360 } 361 } 362 Err(e) => { 363 error!("[IMPORTER] watch error: {:?}", e); 364 } 365 } 366 } 367 368 pub(crate) async fn listen_for_files( 369 &self, 370 walk_import_folders: bool, 371 walk_export_folders: bool, 372 ) -> color_eyre::Result<Debouncer<RecommendedWatcher, FileIdMap>> { 373 // Create a file listener using notify to listen for new files to import 374 let jetstream = self.jetstream.clone(); 375 let tokio_runtime = tokio::runtime::Handle::current(); 376 let debouncer = new_debouncer( 377 Duration::from_secs(2), 378 None, 379 move |res: DebounceEventResult| { 380 Importer::file_event_handler(res, jetstream.clone(), tokio_runtime.clone()); 381 }, 382 )?; 383 384 if walk_import_folders { 385 warn!("Walking input dir to see if we missed any. Errors on missing files on the consumer are expected"); 386 walkdir::WalkDir::new(&self.import_folder) 387 .into_iter() 388 .filter_map(|e| e.ok()) 389 .filter(|e| e.file_type().is_file()) 390 .filter(|e| { 391 let extension = e 392 .path() 393 .extension() 394 .unwrap_or_default() 395 .to_string_lossy() 396 .to_lowercase(); 397 ALLOWED_TYPES.contains(&extension.as_str()) 398 }) 399 .for_each(|entry| { 400 let path = entry.path().to_path_buf(); 401 let jetstream = self.jetstream.clone(); 402 tokio::spawn(async move { 403 Importer::file_added(path, jetstream).await.unwrap(); 404 tokio::time::sleep(std::time::Duration::from_millis(100)).await; 405 }); 406 }); 407 warn!("Done walking input dir"); 408 } 409 if walk_export_folders { 410 warn!("Walking output dir to ensure signing worked"); 411 walkdir::WalkDir::new(&self.export_folder) 412 .into_iter() 413 .filter_map(|e| e.ok()) 414 .filter(|e| e.file_type().is_file()) 415 .filter(|e| { 416 let extension = e 417 .path() 418 .extension() 419 .unwrap_or_default() 420 .to_string_lossy() 421 .to_lowercase(); 422 ALLOWED_TYPES.contains(&extension.as_str()) 423 }) 424 // Filter any which already have a c2pa file with the same name next to it 425 .filter(|e| { 426 let path = e.path(); 427 let sidecar_path = path.with_extension("c2pa"); 428 !sidecar_path.exists() 429 }) 430 .for_each(|entry| { 431 let path = entry.path().to_path_buf(); 432 let jetstream = self.jetstream.clone(); 433 let mut exif_data = None; 434 // Get exif data 435 { 436 let meta = rexiv2::Metadata::new_from_path(&path); 437 match meta { 438 Ok(exif) => { 439 trace!("Got exif data"); 440 441 let exif = self.extract_exif(exif); 442 exif_data = Some(exif); 443 444 debug!("Finished processing exif data") 445 } 446 Err(e) => { 447 error!("[IMPORTER] error reading exif data: {:?}", e); 448 } 449 } 450 }; 451 tokio::spawn(async move { 452 // Add a c2pa event to the stream 453 let c2pa_event = super::jetstream_events::C2PASignRequest { 454 path: path.canonicalize().unwrap(), 455 exif: exif_data, 456 }; 457 let bytes = bincode::serialize(&c2pa_event).unwrap(); 458 publish_with_error_printing( 459 &jetstream, 460 "import.sign_request".to_string(), 461 bytes, 462 ) 463 .await; 464 tokio::time::sleep(Duration::from_millis(100)).await; 465 }); 466 }); 467 } 468 Ok(debouncer) 469 } 470 } 471 472 async fn publish_with_error_printing(ctx: &Context, subject: String, data: Vec<u8>) { 473 match ctx.publish(subject.clone(), data.into()).await { 474 Ok(_) => {} 475 Err(e) => { 476 error!( 477 "[IMPORTER] error publishing message to `{}`: {:?}", 478 subject, e 479 ); 480 // Crash the program if we can't publish the message 481 std::process::exit(1); 482 } 483 } 484 }