diff --git a/src/main.rs b/src/main.rs index d1e789d71a..ebe8c485fb 100644 --- a/src/main.rs +++ b/src/main.rs @@ -12,178 +12,178 @@ use std::{fs, io, env}; use std::str::FromStr; struct BucketDownloader { - bucket: Bucket, - destination_directory: PathBuf, - temp_directory: PathBuf, + bucket: Bucket, + destination_directory: PathBuf, + temp_directory: PathBuf, } impl BucketDownloader { - pub fn create( - bucket: Bucket, - destination_directory: &str, - temp_directory: &str) -> AnyhowResult - { - let destination_directory = PathBuf::from(destination_directory); - let temp_directory = PathBuf::from(temp_directory); + pub fn create( + bucket: Bucket, + destination_directory: &str, + temp_directory: &str) -> AnyhowResult + { + let destination_directory = PathBuf::from(destination_directory); + let temp_directory = PathBuf::from(temp_directory); - Self::check_directory(&destination_directory)?; - Self::check_directory(&temp_directory)?; + Self::check_directory(&destination_directory)?; + Self::check_directory(&temp_directory)?; - Ok(Self { - bucket, + Ok(Self { + bucket, destination_directory, - temp_directory, - }) - } + temp_directory, + }) + } - fn check_directory(path: &PathBuf) -> AnyhowResult<()> { - if !path.exists() { - bail!("Path doesn't exist: {:?}", path); - } - if !path.is_dir() { - bail!("Path isn't a directory: {:?}", path); - } + fn check_directory(path: &PathBuf) -> AnyhowResult<()> { + if !path.exists() { + bail!("Path doesn't exist: {:?}", path); + } + if !path.is_dir() { + bail!("Path isn't a directory: {:?}", path); + } Ok(()) - } + } - pub fn download_matching(&self, path: &str) -> AnyhowResult<()> { - let results = self.bucket.list_blocking("".to_string(), None)?; + pub fn download_matching(&self, path: &str) -> AnyhowResult<()> { + let results = self.bucket.list_blocking("".to_string(), None)?; - let mut objects = Vec::new(); + let mut objects = Vec::new(); - for (list, code) in results { + for (list, code) in results { if code != 200 { - bail!("Response from S3 bucket was not okay: {}", code); - } + bail!("Response from S3 bucket was not okay: {}", code); + } - if list.is_truncated { - // TODO: Handle pagination - warn!("Results returned are paginated. This isn't handled yet!"); - } + if list.is_truncated { + // TODO: Handle pagination + warn!("Results returned are paginated. This isn't handled yet!"); + } - for content in list.contents { - objects.push(content.clone()); - } - } + for content in list.contents { + objects.push(content.clone()); + } + } - for object in objects.iter() { - info!("all objects: {:?}", object.key); - } + for object in objects.iter() { + info!("all objects: {:?}", object.key); + } - let objects : Vec = objects.into_iter() - .filter(|o| o.key.starts_with(path)) - .filter(|o| o.size > 0) - .collect(); + let objects : Vec = objects.into_iter() + .filter(|o| o.key.starts_with(path)) + .filter(|o| o.size > 0) + .collect(); - for object in objects.iter() { - self.download_object(object)?; - } + for object in objects.iter() { + self.download_object(object)?; + } - Ok(()) - } + Ok(()) + } fn download_object(&self, object: &Object) -> AnyhowResult<()> { if self.object_already_downloaded(object) { - info!("Object already downloaded: {}", object.key); - return Ok(()); - } - self.download_object_to_temp(object)?; - self.move_temp_to_downloads(object)?; + info!("Object already downloaded: {}", object.key); + return Ok(()); + } + self.download_object_to_temp(object)?; + self.move_temp_to_downloads(object)?; Ok(()) - } + } - fn object_already_downloaded(&self, object: &Object) -> bool { - let mut path = object.key.clone(); - if path.starts_with("/") { - path = path[1..].into(); - } + fn object_already_downloaded(&self, object: &Object) -> bool { + let mut path = object.key.clone(); + if path.starts_with("/") { + path = path[1..].into(); + } - let download_path = self.destination_directory.join(path); - if !download_path.exists() { - info!("Download path does not already exist: {:?}", download_path); - return false; - } + let download_path = self.destination_directory.join(path); + if !download_path.exists() { + info!("Download path does not already exist: {:?}", download_path); + return false; + } - let file_hash = Self::hash_file_contents(&download_path).unwrap(); // TODO - let object_hash = Self::hash_object(object); + let file_hash = Self::hash_file_contents(&download_path).unwrap(); // TODO + let object_hash = Self::hash_object(object); - // NB: Multipart-uploaded objects do not have an accurate md5 hash. + // NB: Multipart-uploaded objects do not have an accurate md5 hash. // The fix for this is to copy them, which recomputes the md5 wholesale. - // Read more: https://stackoverflow.com/a/29350548 - if object_hash.contains("-") { - warn!("Multipart hash '{}' for file {}. This will continue to download indefinitely until fixed.", - object_hash, object.key); - } + // Read more: https://stackoverflow.com/a/29350548 + if object_hash.contains("-") { + warn!("Multipart hash '{}' for file {}. This will continue to download indefinitely until fixed.", + object_hash, object.key); + } - if !file_hash.eq_ignore_ascii_case(&object_hash) { - info!("File hash '{}' does not match object hash '{}' for file {}", - file_hash, object_hash, object.key); - return false; - } + if !file_hash.eq_ignore_ascii_case(&object_hash) { + info!("File hash '{}' does not match object hash '{}' for file {}", + file_hash, object_hash, object.key); + return false; + } - true - } + true + } - fn download_object_to_temp(&self, object: &Object) -> AnyhowResult<()> { - info!("Downloading object: {}", object.key); + fn download_object_to_temp(&self, object: &Object) -> AnyhowResult<()> { + info!("Downloading object: {}", object.key); - let temp_file_path = self.temp_download_path(object); - let mut temp_file = File::create(temp_file_path)?; + let temp_file_path = self.temp_download_path(object); + let mut temp_file = File::create(temp_file_path)?; - let code = self.bucket.get_object_stream_blocking(&object.key, &mut temp_file)?; - if code != 200 { - bail!("Couldn't download object. Code = {}", code); - } + let code = self.bucket.get_object_stream_blocking(&object.key, &mut temp_file)?; + if code != 200 { + bail!("Couldn't download object. Code = {}", code); + } - Ok(()) - } + Ok(()) + } - fn move_temp_to_downloads(&self, object: &Object) -> AnyhowResult<()> { - let full_download_path = self.permanent_download_path(object); + fn move_temp_to_downloads(&self, object: &Object) -> AnyhowResult<()> { + let full_download_path = self.permanent_download_path(object); let download_directory_base = Self::base_directory(&full_download_path); if !download_directory_base.exists() { - info!("Creating directory: {:?}", download_directory_base); - fs::create_dir_all(download_directory_base)?; - } + info!("Creating directory: {:?}", download_directory_base); + fs::create_dir_all(download_directory_base)?; + } - let temp_file_path = self.temp_download_path(object); + let temp_file_path = self.temp_download_path(object); - fs::rename(&temp_file_path, full_download_path)?; + fs::rename(&temp_file_path, full_download_path)?; Ok(()) - } + } - fn temp_download_path(&self, object: &Object) -> PathBuf { - let object_hash = Self::hash_object(object); - self.temp_directory.join(&object_hash) - } + fn temp_download_path(&self, object: &Object) -> PathBuf { + let object_hash = Self::hash_object(object); + self.temp_directory.join(&object_hash) + } - fn permanent_download_path(&self, object: &Object) -> PathBuf { - self.destination_directory.join(&object.key) - } + fn permanent_download_path(&self, object: &Object) -> PathBuf { + self.destination_directory.join(&object.key) + } - fn base_directory(path: &Path) -> PathBuf { - let mut path = path.to_path_buf(); - path.pop(); - path - } + fn base_directory(path: &Path) -> PathBuf { + let mut path = path.to_path_buf(); + path.pop(); + path + } - fn hash_file_contents(path: &Path) -> AnyhowResult { - let mut file = fs::File::open(path)?; - let mut hasher = Md5::new(); - let _n = io::copy(&mut file, &mut hasher)?; - let hash = hasher.result(); + fn hash_file_contents(path: &Path) -> AnyhowResult { + let mut file = fs::File::open(path)?; + let mut hasher = Md5::new(); + let _n = io::copy(&mut file, &mut hasher)?; + let hash = hasher.result(); let hex_encoded = hex::encode(hash); Ok(hex_encoded) - } + } - fn hash_object(object: &Object) -> String { - // NB: E-tag contains quotes, so we remove them - object.e_tag - .trim_start_matches("\"") - .trim_end_matches("\"") - .to_string() - } + fn hash_object(object: &Object) -> String { + // NB: E-tag contains quotes, so we remove them + object.e_tag + .trim_start_matches("\"") + .trim_end_matches("\"") + .to_string() + } } // TODO: Some of these should be yaml/toml configs @@ -196,42 +196,42 @@ const ENV_TEMP_DIR : &'static str = "TEMP_DIR"; const ENV_MATCH_PATH : &'static str = "MATCH_PATH"; fn get_env(env_name: &str) -> AnyhowResult { - match env::var(env_name).as_ref().ok() { - None => bail!("Must set env var {}", env_name), - Some(s) => Ok(s.to_string()), - } + match env::var(env_name).as_ref().ok() { + None => bail!("Must set env var {}", env_name), + Some(s) => Ok(s.to_string()), + } } pub fn main() -> AnyhowResult<()> { - std::env::set_var("RUST_LOG", "info"); - env_logger::init(); + std::env::set_var("RUST_LOG", "info"); + env_logger::init(); - info!("starting"); + info!("starting"); - let access_key = get_env(ENV_ACCESS_KEY)?; - let secret_key = get_env(ENV_SECRET_KEY)?; - let region = get_env(ENV_REGION)?; - let bucket_name = get_env(ENV_BUCKET_NAME)?; - let download_dir = get_env(ENV_DOWNLOAD_DIR)?; - let temp_dir = get_env(ENV_TEMP_DIR)?; - let match_path = get_env(ENV_MATCH_PATH)?; + let access_key = get_env(ENV_ACCESS_KEY)?; + let secret_key = get_env(ENV_SECRET_KEY)?; + let region = get_env(ENV_REGION)?; + let bucket_name = get_env(ENV_BUCKET_NAME)?; + let download_dir = get_env(ENV_DOWNLOAD_DIR)?; + let temp_dir = get_env(ENV_TEMP_DIR)?; + let match_path = get_env(ENV_MATCH_PATH)?; - let credentials = Credentials::new_blocking( - Some(&access_key), - Some(&secret_key), - None, - None, - None, - )?; + let credentials = Credentials::new_blocking( + Some(&access_key), + Some(&secret_key), + None, + None, + None, + )?; - let region = Region::from_str(®ion)?; - let bucket = Bucket::new(&bucket_name, region, credentials)?; + let region = Region::from_str(®ion)?; + let bucket = Bucket::new(&bucket_name, region, credentials)?; - let bucket_downloader = BucketDownloader::create(bucket, &download_dir, &temp_dir)?; + let bucket_downloader = BucketDownloader::create(bucket, &download_dir, &temp_dir)?; - bucket_downloader.download_matching(&match_path)?; + bucket_downloader.download_matching(&match_path)?; - info!("Done"); + info!("Done"); - Ok(()) + Ok(()) }