photosync/src/main.rs

use std::collections::HashMap;
use std::error::Error;
use std::fs;
use std::fs::File;
use std::io;
use std::net::SocketAddr;
use std::path::PathBuf;

use google_api_auth;
use google_photoslibrary1 as photos;
use hexihasher;
use lazy_static::lazy_static;
use log::{debug, info};
use photos::schemas::{Album, MediaItem, SearchMediaItemsRequest};
use regex::Regex;
use reqwest;
use structopt::StructOpt;
use yup_oauth2::{Authenticator, InstalledFlow};

mod web;

#[derive(Debug, StructOpt)]
enum Command {
    /// List albums for the user of the given credentials.  Optionally title filter.
    ListAlbums {
        title_filter: Option<Regex>,
    },
    SearchMediaItems {
        album_id: String,
    },
    Sync {
        /// Optional album title to filter.  Default will mirror all albums.
        #[structopt(short, long)]
        title_filter: Option<Regex>,
        /// Directory to store sync.
        output: PathBuf,
    },
    Serve {
        /// Directory of data fetched by `sync`.
        root: PathBuf,
        /// HTTP address to listen for web requests.
        #[structopt(long = "addr", default_value = "0.0.0.0:0")]
        addr: SocketAddr,
    },
}

#[derive(Debug, StructOpt)]
#[structopt(
    name = "photosync",
    about = "Utility for interacting with Google Photos API."
)]
struct Opt {
    /// Activate debug mode
    #[structopt(short, parse(from_occurrences))]
    verbose: usize,

    /// Path to json file containing Google client ID and secrets for out of band auth flow.
    #[structopt(long)]
    credentials: PathBuf,
    /// Path to json file where photosync will store auth tokens refreshed from Google.
    #[structopt(long)]
    token_cache: PathBuf,

    #[structopt(subcommand)]
    cmd: Command,
}

fn new_client(
    credentials: &PathBuf,
    token_cache: &PathBuf,
) -> Result<photos::Client, Box<dyn Error>> {
    let secret = yup_oauth2::read_application_secret(credentials)?;

    // Create an authenticator that uses an InstalledFlow to authenticate. The
    // authentication tokens are persisted to a file named tokencache.json. The
    // authenticator takes care of caching tokens to disk and refreshing tokens once
    // they've expired.
    let auth = Authenticator::new(InstalledFlow::new(
        secret,
        yup_oauth2::InstalledFlowReturnMethod::Interactive,
    ))
    .persist_tokens_to_disk(token_cache)
    .build()
    .unwrap();

    let scopes = vec![photos::scopes::PHOTOSLIBRARY_READONLY];

    let auth = google_api_auth::yup_oauth2::from_authenticator(auth, scopes);

    Ok(photos::Client::new(auth))
}

struct SearchIter<'a> {
    client: &'a photos::Client,
    items: ::std::vec::IntoIter<MediaItem>,
    finished: bool,
    req: SearchMediaItemsRequest,
}

impl<'a> SearchIter<'a> {
    fn new(client: &'a photos::Client, req: SearchMediaItemsRequest) -> Self {
        SearchIter {
            client,
            items: Vec::new().into_iter(),
            finished: false,
            req,
        }
    }
}

impl<'a> Iterator for SearchIter<'a> {
    type Item = Result<MediaItem, photos::Error>;

    fn next(&mut self) -> Option<Result<MediaItem, photos::Error>> {
        loop {
            if let Some(v) = self.items.next() {
                return Some(Ok(v));
            }
            if self.finished {
                return None;
            }
            let resp = match self
                .client
                .media_items()
                .search(self.req.clone())
                .execute_with_default_fields()
            {
                Ok(resp) => resp,
                Err(err) => return Some(Err(err)),
            };
            if resp.next_page_token.is_none() {
                self.finished = true;
            }
            self.req.page_token = resp.next_page_token;
            if let Some(items) = resp.media_items {
                self.items = items.into_iter();
            }
        }
    }
}

fn print_media_items(media_items: Vec<MediaItem>) {
    for mi in &media_items {
        let id = mi
            .id
            .as_ref()
            .map_or("NO ID".to_string(), |s| s.to_string());
        println!(
            "media item: {}\n\t{}\n\t{}",
            mi.filename.as_ref().unwrap_or(&"NO FILENAME".to_string()),
            hexihasher::sha256(id.as_bytes()),
            id,
        );
    }
    println!("({}) items total", media_items.len());
}

fn search_media_items(
    client: &photos::Client,
    album_id: &str,
) -> Result<Vec<MediaItem>, Box<dyn Error>> {
    let media_items = SearchIter::new(
        &client,
        SearchMediaItemsRequest {
            album_id: Some(album_id.to_string()),
            // 100 is the documented max.
            page_size: Some(100),
            ..Default::default()
        },
    )
    .filter_map(|mi| mi.ok())
    .collect();
    Ok(media_items)
}

lazy_static! {
    static ref MIME_TO_EXT: HashMap<&'static str, &'static str> = [
        ("image/gif", "gif"),
        ("image/heif", "heic"),
        ("image/jpeg", "jpg"),
    ]
    .iter()
    .copied()
    .collect();
}

fn sync_albums(
    client: &photos::Client,
    title_filter: Option<Regex>,
    output_dir: PathBuf,
) -> Result<(), Box<dyn Error>> {
    // Put images from all albums in common directory.
    let image_dir = output_dir.join("images");
    if !image_dir.exists() {
        fs::create_dir_all(&image_dir)?;
    }
    let albums = list_albums(client, title_filter)?;
    for a in &albums {
        let album_id = a.id.as_ref().expect("unset album id").to_string();
        let album_dir = output_dir.join(&album_id);
        if !album_dir.exists() {
            info!("making album directory {}", album_dir.to_string_lossy());
            fs::create_dir_all(&album_dir)?;
        }

        let album = search_media_items(client, &album_id)?;
        for (i, mi) in album.iter().enumerate() {
            let mi_id = mi.id.as_ref().expect("unset media item id").to_string();
            let filename = mi
                .filename
                .as_ref()
                .map_or("NO_FILENAME".to_string(), |s| s.to_string());
            let image_path = image_dir.join(mi_id);
            if image_path.exists() {
                info!(
                    "Skipping already downloaded {} @ {}",
                    &filename,
                    image_path.to_string_lossy()
                );
            } else {
                let download_path = image_path.with_extension("download");
                info!(
                    "({}/{}) Downloading {} -> {}",
                    i + 1,
                    &album.len(),
                    &filename,
                    download_path.to_string_lossy()
                );
                let base_url = mi.base_url.as_ref().expect("missing base_url");
                let url = format!("{}=d", base_url);
                let mut r = reqwest::blocking::get(&url)?;
                let mut w = File::create(&download_path)?;
                let _n = io::copy(&mut r, &mut w)?;
                info!(
                    "Rename {} -> {}",
                    download_path.to_string_lossy(),
                    image_path.to_string_lossy()
                );
                fs::rename(download_path, &image_path)?;
            }
        }
        let j = serde_json::to_string(&album)?;
        let path = album_dir.join("album.json");
        info!("saving {}", path.to_string_lossy());
        fs::write(path, j)?;
    }
    // Serialize it to a JSON string.
    let j = serde_json::to_string(&albums)?;

    let path = output_dir.join("albums.json");
    info!("saving {}", path.to_string_lossy());
    fs::write(path, j)?;
    Ok(())
}

fn print_albums(albums: Vec<Album>) {
    for a in albums {
        println!(
            "album: {} {} ({} items)",
            a.id.unwrap_or("NO ID".to_string()),
            a.title.unwrap_or("NO TITLE".to_string()).to_string(),
            a.media_items_count.unwrap_or(0)
        );
    }
}

fn list_albums(
    client: &photos::Client,
    title_filter: Option<Regex>,
) -> Result<Vec<Album>, Box<dyn Error>> {
    Ok(client
        .shared_albums()
        .list()
        .iter_shared_albums_with_all_fields()
        .filter_map(|a| a.ok())
        .filter(|a| {
            match (&title_filter, &a.title) {
                // keep everything when no filter or title.
                (None, None) => true,
                // skip when filter given but the media item doesn't have a title (it can't match)
                (_, None) => false,
                // skip when the media item doesn't match the filter
                (Some(title_filter), Some(title)) if !title_filter.is_match(&title) => false,
                // keep everything else
                _ => true,
            }
        })
        .collect())
}

pub fn serve(addr: SocketAddr, root: PathBuf) -> Result<(), Box<dyn Error>> {
    web::run(addr, root)
}

fn main() -> Result<(), Box<dyn std::error::Error>> {
    let opt = Opt::from_args();
    stderrlog::new()
        .module(module_path!())
        .verbosity(opt.verbose)
        .init()
        .unwrap();
    debug!("opt: {:?}", opt);
    let client = new_client(&opt.credentials, &opt.token_cache)?;
    match opt.cmd {
        Command::ListAlbums { title_filter } => {
            print_albums(list_albums(&client, title_filter)?);
            Ok(())
        }
        Command::SearchMediaItems { album_id } => {
            print_media_items(search_media_items(&client, &album_id)?);
            Ok(())
        }
        Command::Sync {
            title_filter,
            output,
        } => sync_albums(&client, title_filter, output),
        Command::Serve { addr, root } => serve(addr, root),
    }
}