photosync/src/main.rs

use std::collections::HashMap;
use std::error::Error;
use std::net::SocketAddr;
use std::path::PathBuf;
use std::sync::{Arc, Mutex};
use std::thread;
use std::time;

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

use photosync::library::Library;
use photosync::rweb;

fn parse_duration(src: &str) -> Result<time::Duration, std::num::ParseIntError> {
    let secs = str::parse::<u64>(src)?;
    Ok(time::Duration::from_secs(secs))
}

#[derive(Debug, StructOpt)]
struct Sync {
    #[structopt(flatten)]
    auth: Auth,
    /// Optional album title to filter.  Default will mirror all albums.
    #[structopt(short, long)]
    title_filter: Option<Regex>,
    /// Directory to store sync.
    root: PathBuf,
}

#[derive(Debug, StructOpt)]
struct 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)]
enum Command {
    /// List albums for the user of the given credentials.  Optionally title filter.
    ListAlbums {
        #[structopt(flatten)]
        auth: Auth,
        title_filter: Option<Regex>,
    },
    SearchMediaItems {
        #[structopt(flatten)]
        auth: Auth,
        album_id: String,
    },
    Sync {
        #[structopt(flatten)]
        sync: Sync,
    },
    Serve {
        #[structopt(flatten)]
        serve: Serve,
    },
    ServeAndSync {
        /// Sync albums at given interval.
        #[structopt(parse(try_from_str = parse_duration))]
        interval: time::Duration,

        #[structopt(flatten)]
        sync: Sync,
        /// HTTP address to listen for web requests.
        #[structopt(long = "addr", default_value = "0.0.0.0:0")]
        addr: SocketAddr,
    },
}

#[derive(Debug, StructOpt)]
struct Auth {
    /// 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,
}

#[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,

    #[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>,
    lib: &Library,
) -> Result<(), Box<dyn Error>> {
    let albums = list_albums(client, title_filter)?;
    info!("albums {:?}", albums);
    lib.create_album_index(&albums)?;
    for a in &albums {
        let album_id = a.id.as_ref().expect("unset album id").to_string();
        let media_items = search_media_items(client, &album_id)?;
        lib.create_album(&album_id, &media_items)?;
        for (i, mi) in media_items.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 base_url = mi.base_url.as_ref().expect("missing base_url");
            let image_path = lib.download_image(&filename, &mi_id, &base_url)?;
            info!(
                "({}/{}) Checking {} -> {}",
                i + 1,
                &media_items.len(),
                &filename,
                image_path.to_string_lossy()
            );
        }
    }
    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
        .albums()
        .list()
        .iter_albums_with_all_fields()
        .chain(
            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())
}

fn background_sync(
    client: photos::Client,
    interval: time::Duration,
    title_filter: Option<Regex>,
    lib: Library,
) -> Result<(), Box<dyn Error>> {
    thread::spawn(move || loop {
        if let Err(err) = sync_albums(&client, &title_filter, &lib) {
            error!("Error syncing: {}", err);
        }
        thread::sleep(interval);
    });
    Ok(())
}

pub fn serve(addr: SocketAddr, lib: Library) -> Result<(), Box<dyn Error>> {
    rweb::run(addr, lib)
}

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 image_cache: Mutex<Box<dyn Cacher>> =
        Mutex::new(Box::new(S3Cacher::new("photosync".to_string())?));
    let image_cache = Arc::new(image_cache);
    match opt.cmd {
        Command::ListAlbums { auth, title_filter } => {
            let client = new_client(&auth.credentials, &auth.token_cache)?;
            print_albums(list_albums(&client, &title_filter)?);
            Ok(())
        }
        Command::SearchMediaItems { auth, album_id } => {
            let client = new_client(&auth.credentials, &auth.token_cache)?;
            print_media_items(search_media_items(&client, &album_id)?);
            Ok(())
        }
        Command::Sync {
            sync:
                Sync {
                    auth,
                    title_filter,
                    root,
                },
        } => {
            let client = new_client(&auth.credentials, &auth.token_cache)?;
            let lib = Library::new(root, image_cache)?;
            sync_albums(&client, &title_filter, &lib)?;
            Ok(())
        }
        Command::Serve {
            serve: Serve { addr, root },
        } => {
            let lib = Library::new(root, image_cache)?;
            serve(addr, lib)
        }
        Command::ServeAndSync {
            interval,
            sync:
                Sync {
                    auth,
                    title_filter,
                    root,
                },
            addr,
        } => {
            let client = new_client(&auth.credentials, &auth.token_cache)?;
            let lib = Library::new(root, image_cache)?;
            background_sync(client, interval, title_filter, lib.clone())?;
            serve(addr, lib)?;
            Ok(())
        }
    }
}