Albatross/src/main.rs

extern crate serde;

#[macro_use]
extern crate serde_derive;

use chrono::{NaiveDateTime, Utc};
use clap::{App, Arg};
use flate2::write::GzEncoder;
use flate2::Compression;
use regex::Regex;
use std::fs::{copy, create_dir, create_dir_all, remove_dir_all, remove_file, DirEntry, File};
use std::path::PathBuf;
use std::time::Instant;

mod albatross_config;

use albatross_config::{AlbatrossConfig, WorldConfig, WorldType};
use std::collections::HashMap;

/// Struct to store information about the region
struct Region {
    /// x position of the region
    x: i64,
    /// y position of the region
    y: i64,
}

impl Region {
    fn from_string(string: String) -> Option<Self> {
        let re = Regex::new(r"r\.(?P<x>-?[0-9]*)+\.(?P<y>-?[0-9]*)").unwrap();
        if re.is_match(string.as_str()) {
            let captures = re.captures(string.as_str()).unwrap();

            return Some(Region {
                x: captures["x"].parse::<i64>().unwrap(),
                y: captures["y"].parse::<i64>().unwrap(),
            });
        }

        None
    }
}

/// Backup a directory
///
/// # Param
/// * `world_path` - path to the world folder
/// * `backup_path` - path to the backup folder
fn backup_dir(
    dir_name: &str,
    world_path: &PathBuf,
    backup_path: &PathBuf,
) -> Result<u64, std::io::Error> {
    let mut src_dir = world_path.clone();
    src_dir.push(dir_name);
    let mut backup_dir = backup_path.clone();
    backup_dir.push(dir_name);
    create_dir(&backup_dir)?;

    let mut file_count = 0;
    for entry in src_dir.read_dir()? {
        let entry = entry?;
        let mut target = backup_dir.clone();
        target.push(entry.file_name());

        copy(entry.path(), target)?;
        file_count += 1;
    }

    Ok(file_count)
}

/// Backup the regions
///
/// # Param
/// * `dir_name` - name of the backup folder
/// * `save_radius` - block radius to save
/// * `world_path` - path to the world folder
/// * `backup_path` - path to the backup folder
fn backup_region(
    dir_name: &str,
    save_radius: u64,
    world_path: &PathBuf,
    backup_path: &PathBuf,
) -> Result<u64, std::io::Error> {
    let mut count: u64 = 0;
    let mut src_dir = world_path.clone();
    src_dir.push(dir_name);
    let mut backup_dir = backup_path.clone();
    backup_dir.push(dir_name);
    create_dir(&backup_dir)?;

    let save_radius = (save_radius as f64 / 512.0).ceil() as i64;

    for entry in src_dir.read_dir()? {
        let entry = entry?;
        let file_name = entry.file_name().to_str().unwrap().to_string();

        if let Some(region) = Region::from_string(file_name) {
            if region.x.abs() <= save_radius && region.y.abs() <= save_radius {
                let mut target = backup_dir.clone();
                target.push(entry.file_name());

                copy(entry.path(), target)?;
                count += 1;
            }
        }
    }

    Ok(count)
}

/// Backup a world
///
/// # Param
/// * `world_path` - path to the world folder
/// * `backup_path` - path to the backup folder
/// * `world_config` - world config options
fn backup_world(
    world_path: PathBuf,
    backup_path: PathBuf,
    world_config: &WorldConfig,
) -> Result<u64, std::io::Error> {
    let mut backup_path = backup_path.clone();
    let region_count;
    backup_path.push(&world_config.world_name);
    create_dir(backup_path.as_path())?;

    backup_region("poi", world_config.save_radius, &world_path, &backup_path)?;
    region_count = backup_region(
        "region",
        world_config.save_radius,
        &world_path,
        &backup_path,
    )?;
    backup_dir("data", &world_path, &backup_path)?;

    Ok(region_count)
}

/// Backup the overworld
///
/// # Param
/// * `world_path` - path to the world folder
/// * `backup_path` - path to the backup folder
/// * `world_config` - world config options
fn backup_overworld(
    world_path: PathBuf,
    backup_path: PathBuf,
    world_config: &WorldConfig,
) -> Result<u64, std::io::Error> {
    backup_world(world_path, backup_path, world_config)
}

/// Backup the nether
///
/// # Param
/// * `world_path` - path to the world folder
/// * `backup_path` - path to the backup folder
/// * `world_config` - world config options
fn backup_nether(
    world_path: PathBuf,
    backup_path: PathBuf,
    world_config: &WorldConfig,
) -> Result<u64, std::io::Error> {
    let mut nether_path = world_path.clone();
    nether_path.push("DIM-1");

    backup_world(nether_path, backup_path, world_config)
}

/// Backup the end
///
/// # Param
/// * `world_path` - path to the world folder
/// * `backup_path` - path to the backup folder
/// * `world_config` - world config options
fn backup_end(
    world_path: PathBuf,
    backup_path: PathBuf,
    world_config: &WorldConfig,
) -> Result<u64, std::io::Error> {
    let mut end_path = world_path.clone();
    end_path.push("DIM1");

    backup_world(end_path, backup_path, world_config)
}

/// Compress the backup after the files have been copied
///
/// # Param
/// * `tmp_dir`: tmp directory with the backed up files
/// * `output_file`: output archive
fn compress_backup(tmp_dir: &PathBuf, output_file: &PathBuf) -> Result<(), std::io::Error> {
    let archive = File::create(output_file)?;
    let enc = GzEncoder::new(archive, Compression::default());
    let mut tar_builder = tar::Builder::new(enc);
    tar_builder.append_dir_all(".", tmp_dir)?;
    Ok(())
}

/// Get the time of the backup from a file name
///
/// # Param
/// * `archive_entry`: archive entry
fn get_time_from_file_name(
    archive_entry: &DirEntry,
) -> Result<Option<NaiveDateTime>, std::io::Error> {
    let file_name = archive_entry.file_name().to_str().unwrap().to_string();
    let name: Vec<&str> = file_name.split("_").collect();

    Ok(chrono::NaiveDateTime::parse_from_str(name[0], "%d-%m-%y-%H:%M:%S").ok())
}

/// Removes the old backups from the ouput directory
///
/// # Params
/// * `output_dir` - output directory containing
/// * `keep` - number of backups to keep
fn remove_old_backups(output_dir: &PathBuf, keep: u64) -> Result<usize, std::io::Error> {
    let mut backups = vec![];
    let mut num_of_removed_backups: usize = 0;

    for entry in output_dir.read_dir()? {
        let entry = entry?;

        if let Some(ext) = entry.path().extension() {
            if ext == "gz" {
                backups.push(entry);
            }
        }
    }

    if backups.len() > keep as usize {
        backups.sort_by(|a, b| {
            let a_time = get_time_from_file_name(a).unwrap().unwrap();
            let b_time = get_time_from_file_name(b).unwrap().unwrap();

            b_time.cmp(&a_time)
        });

        num_of_removed_backups = backups.len() - keep as usize;

        for _i in 0..num_of_removed_backups {
            let oldest = backups.pop().unwrap();
            remove_file(oldest.path())?;
        }
    }

    Ok(num_of_removed_backups)
}

/// Sends a webhook to Discord if its configured
///
/// # Params
/// * `msg` - Message to send to discord
/// * `cfg` - Albatross config
fn send_webhook(msg: &str, cfg: &AlbatrossConfig) {
    if let Some(webhook) = &cfg.backup.discord_webhook {
        let mut map = HashMap::new();

        map.insert("content".to_string(), msg.to_string());

        let client = reqwest::blocking::Client::new();
        client.post(webhook).json(&map).send().ok();
    }
}

/// Backup the configured worlds from a minecraft server
///
/// # Params
/// * `cfg` - config file
fn do_backup(cfg: AlbatrossConfig) -> Result<(), std::io::Error> {
    let server_base_dir = cfg.backup.minecraft_dir.clone();
    let worlds = cfg.world_config.clone().expect("No worlds configured");
    let time_str = Utc::now().format("%d-%m-%y-%H:%M:%S").to_string();
    let backup_name = format!("{}_backup.tar.gz", time_str);
    let mut output_archive = cfg.backup.output_dir.clone();
    output_archive.push(backup_name);
    let mut tmp_dir = cfg.backup.output_dir.clone();
    tmp_dir.push("tmp");
    remove_dir_all(&tmp_dir).ok();

    create_dir_all(tmp_dir.clone()).unwrap();

    send_webhook("**Albatross is swooping in to backup your worlds!**", &cfg);
    let timer = Instant::now();
    for world in worlds {
        let mut world_dir = server_base_dir.clone();
        let world_name = world.world_name.clone();
        let world_type = match world.world_type.clone() {
            Some(world_type) => world_type,
            None => WorldType::OVERWORLD,
        };
        world_dir.push(world_name.clone());

        if world_dir.exists() && world_dir.is_dir() {
            send_webhook(
                format!("Starting backup of **{}**", world_name).as_str(),
                &cfg,
            );
            match world_type {
                WorldType::OVERWORLD => {
                    let region_count =
                        backup_overworld(world_dir.clone(), tmp_dir.clone(), &world)?;
                    let player_count = backup_dir("playerdata", &world_dir, &tmp_dir)?;
                    send_webhook(
                        format!(
                            "{} regions and {} player files backed up.",
                            region_count, player_count
                        )
                        .as_str(),
                        &cfg,
                    );
                }
                WorldType::NETHER => {
                    let region_count = backup_nether(world_dir, tmp_dir.clone(), &world)?;
                    send_webhook(
                        format!("{} regions backed up.", region_count).as_str(),
                        &cfg,
                    );
                }
                WorldType::END => {
                    let region_count = backup_end(world_dir, tmp_dir.clone(), &world)?;
                    send_webhook(
                        format!("{} regions backed up.", region_count).as_str(),
                        &cfg,
                    );
                }
            };
        } else {
            send_webhook(format!("Error: {} not found.", world_name).as_str(), &cfg);
            println!("World \"{}\" not found", world_name.clone());
        }
    }

    compress_backup(&tmp_dir, &output_archive)?;

    remove_dir_all(&tmp_dir)?;

    let backups_removed = remove_old_backups(&cfg.backup.output_dir, cfg.backup.backups_to_keep)?;

    if backups_removed > 0 {
        let msg = format!(
            "Albatross mistook **{}** of your old backups for some french fries and ate them!! SKRAWWWW",
            backups_removed
        );
        send_webhook(msg.as_str(), &cfg);
    }

    let secs = timer.elapsed().as_secs();
    send_webhook(
        format!("**Full backup completed in {}s**! *SKREEEEEEEEEE*", secs).as_str(),
        &cfg,
    );
    Ok(())
}

/// Albatross
///
/// Run backups of a Minecraft world
fn main() {
    let mut app = App::new("Albatross").about("Backup your worlds").arg(
        Arg::with_name("config")
            .index(1)
            .short("c")
            .long("config")
            .value_name("CONFIG_PATH")
            .help("Config file path"),
    );
    // Get arg parser
    let matches = app.clone().get_matches();

    if let Some(cfg_path) = matches.value_of("config") {
        let cfg = AlbatrossConfig::new(cfg_path).expect("Config not found");

        if cfg.world_config.is_some() {
            println!("Starting backup");
            match do_backup(cfg) {
                Err(e) => println!("Error doing backup: {}", e),
                _ => {}
            }

            println!("Backup complete");
        } else {
            println!("No worlds specified to backed up!")
        }
    } else {
        app.print_help().expect("Unable to print help");
    }
}