modmapper/src/main.rs

use anyhow::Result;
use argh::FromArgs;
use chrono::{NaiveDateTime, NaiveTime};
use dotenv::dotenv;
use humansize::{file_size_opts, FileSize};
use models::file::File;
use models::game_mod::Mod;
use reqwest::StatusCode;
use sqlx::postgres::PgPoolOptions;
use std::collections::{HashMap, HashSet};
use std::env;
use std::io::Seek;
use std::io::SeekFrom;
use std::io::Write;
use std::process::Command;
use std::time::Duration;
use tempfile::tempdir;
use tokio::io::{AsyncReadExt, AsyncSeekExt};
use tokio::time::sleep;
use tracing::{debug, error, info, info_span, warn};
use unrar::Archive;
use walkdir::WalkDir;

mod extractors;
mod models;
mod nexus_api;
mod nexus_scraper;
mod plugin_processor;

use models::cell;
use models::file;
use models::game;
use models::{game_mod, game_mod::UnsavedMod};
use nexus_api::{GAME_ID, GAME_NAME};
use plugin_processor::process_plugin;

const REQUEST_TIMEOUT: Duration = Duration::from_secs(7200); // 2 hours
const CONNECT_TIMEOUT: Duration = Duration::from_secs(30);

#[derive(FromArgs)]
/// Downloads every mod off nexus mods, parses CELL and WRLD data from plugins in each, and saves the data to the database.
struct Args {
    #[argh(option, short = 'p', default = "1")]
    /// the page number to start scraping for mods on nexus mods.
    page: usize,

    /// file to output the cell mod edit counts as json
    #[argh(option, short = 'e')]
    dump_edits: Option<String>,

    /// folder to output all cell data as json files
    #[argh(option, short = 'c')]
    cell_data: Option<String>,

    /// backfill mods data from the Nexus API (temporary)
    #[argh(switch, short = 'b')]
    backfill_mods: bool,
}

async fn extract_with_compress_tools(
    file: &mut std::fs::File,
    pool: &sqlx::Pool<sqlx::Postgres>,
    db_file: &File,
    db_mod: &Mod,
) -> Result<()> {
    let extractor = extractors::compress_tools::Extractor::new(file);
    for plugin in extractor.into_iter() {
        let (file_path, mut plugin_buf) = plugin?;
        let plugin_span = info_span!("plugin", name = ?file_path);
        let _plugin_span = plugin_span.enter();
        process_plugin(&mut plugin_buf, &pool, &db_file, &db_mod, &file_path).await?;
    }
    Ok(())
}

async fn extract_with_7zip(
    file: &mut std::fs::File,
    pool: &sqlx::Pool<sqlx::Postgres>,
    db_file: &File,
    db_mod: &Mod,
) -> Result<()> {
    file.seek(SeekFrom::Start(0))?;
    let temp_dir = tempdir()?;
    let temp_file_path = temp_dir.path().join("download.zip");
    let mut temp_file = std::fs::File::create(&temp_file_path)?;
    std::io::copy(file, &mut temp_file)?;
    drop(temp_file); // close handle to temp file so 7zip process can open it
    let extracted_path = temp_dir.path().join("extracted");

    Command::new("7z")
        .args(&[
            "x",
            &format!("-o{}", &extracted_path.to_string_lossy()),
            &temp_file_path.to_string_lossy().to_string(),
        ])
        .status()?;

    for entry in WalkDir::new(&extracted_path)
        .contents_first(true)
        .into_iter()
        .filter_entry(|e| {
            if let Some(extension) = e.path().extension() {
                extension == "esp" || extension == "esm" || extension == "esl"
            } else {
                false
            }
        })
    {
        let entry = entry?;
        let file_path = entry.path();
        let plugin_span = info_span!("plugin", name = ?file_path);
        let _plugin_span = plugin_span.enter();
        info!("processing uncompressed file from downloaded archive");
        let mut plugin_buf = std::fs::read(extracted_path.join(file_path))?;
        process_plugin(
            &mut plugin_buf,
            &pool,
            &db_file,
            &db_mod,
            &file_path.to_string_lossy(),
        )
        .await?;
    }
    Ok(())
}

async fn extract_with_unrar(
    file: &mut std::fs::File,
    pool: &sqlx::Pool<sqlx::Postgres>,
    db_file: &File,
    db_mod: &Mod,
    checked_metadata: bool,
) -> Result<()> {
    let temp_dir = tempdir()?;
    let temp_file_path = temp_dir.path().join("download.rar");
    let mut temp_file = std::fs::File::create(&temp_file_path)?;
    std::io::copy(file, &mut temp_file)?;

    let mut plugin_file_paths = Vec::new();
    let list = Archive::new(&temp_file_path.to_string_lossy().to_string())?.list();
    match list {
        Ok(list) => {
            for entry in list.flatten() {
                if let Some(extension) = entry.filename.extension() {
                    if entry.is_file()
                        && (extension == "esp" || extension == "esm" || extension == "esl")
                    {
                        plugin_file_paths.push(entry.filename);
                    }
                }
            }
        }
        Err(_) => {
            if !checked_metadata {
                warn!("failed to read archive and server has no metadata, skipping file");
                file::update_unable_to_extract_plugins(&pool, db_file.id, true).await?;
                return Ok(());
            } else {
                error!("failed to read archive, but server had metadata");
                panic!("failed to read archive, but server had metadata");
            }
        }
    }
    info!(
        num_plugin_files = plugin_file_paths.len(),
        "listed plugins in downloaded archive"
    );

    if !plugin_file_paths.is_empty() {
        info!("uncompressing downloaded archive");
        let extract = Archive::new(&temp_file_path.to_string_lossy().to_string())?
            .extract_to(temp_dir.path().to_string_lossy().to_string());

        let mut extract = match extract {
            Err(err) => {
                warn!(error = %err, "failed to extract with unrar");
                file::update_unable_to_extract_plugins(&pool, db_file.id, true).await?;
                return Ok(());
            }
            Ok(extract) => extract,
        };
        if let Err(err) = extract.process() {
            warn!(error = %err, "failed to extract with unrar");
            file::update_unable_to_extract_plugins(&pool, db_file.id, true).await?;
            return Ok(());
        }

        for file_path in plugin_file_paths.iter() {
            info!(
                ?file_path,
                "processing uncompressed file from downloaded archive"
            );
            let mut plugin_buf = std::fs::read(temp_dir.path().join(file_path))?;
            process_plugin(
                &mut plugin_buf,
                &pool,
                &db_file,
                &db_mod,
                &file_path.to_string_lossy(),
            )
            .await?;
        }
    }
    Ok(())
}

#[tokio::main]
pub async fn main() -> Result<()> {
    dotenv().ok();

    tracing_subscriber::fmt::init();

    let pool = PgPoolOptions::new()
        .max_connections(5)
        .connect(&env::var("DATABASE_URL")?)
        .await?;

    let args: Args = argh::from_env();

    let client = reqwest::Client::builder()
        .timeout(REQUEST_TIMEOUT)
        .connect_timeout(CONNECT_TIMEOUT)
        .build()?;

    if args.backfill_mods {
        for game_mod in game_mod::bulk_get_need_backfill(&pool).await? {
            let response = nexus_api::game_mod::get(&client, game_mod.nexus_mod_id).await?;
            let mod_data = response.extract_data()?;
            game_mod::update_from_api_response(&pool, &game_mod, &mod_data).await?;
            info!(
                id = game_mod.id,
                nexus_mod_id = game_mod.nexus_mod_id,
                "backfilled mod data from api"
            );
        }
    }

    if let Some(dump_edits) = args.dump_edits {
        let mut cell_mod_edit_counts = HashMap::new();
        for x in -77..75 {
            for y in -50..44 {
                if let Some(count) = cell::count_mod_edits(&pool, "Skyrim.esm", 1, x, y).await? {
                    cell_mod_edit_counts.insert(format!("{},{}", x, y), count);
                }
            }
        }
        let mut file = std::fs::File::create(dump_edits)?;
        write!(file, "{}", serde_json::to_string(&cell_mod_edit_counts)?)?;
        return Ok(());
    }

    if let Some(cell_data_dir) = args.cell_data {
        for x in -77..75 {
            for y in -50..44 {
                if let Ok(data) = cell::get_cell_data(&pool, "Skyrim.esm", 1, x, y).await {
                    let path = format!("{}/{}", &cell_data_dir, x);
                    let path = std::path::Path::new(&path);
                    std::fs::create_dir_all(&path)?;
                    let path = path.join(format!("{}.json", y));
                    let mut file = std::fs::File::create(path)?;
                    write!(file, "{}", serde_json::to_string(&data)?)?;
                }
            }
        }
        return Ok(());
    }

    let mut page = args.page;
    let mut has_next_page = true;

    let game = game::insert(&pool, GAME_NAME, GAME_ID as i32).await?;

    while has_next_page {
        let page_span = info_span!("page", page);
        let _page_span = page_span.enter();
        let mod_list_resp = nexus_scraper::get_mod_list_page(&client, page).await?;
        let scraped = mod_list_resp.scrape_mods()?;

        has_next_page = scraped.has_next_page;
        let processed_mods = game_mod::bulk_get_last_updated_by_nexus_mod_ids(
            &pool,
            &scraped
                .mods
                .iter()
                .map(|scraped_mod| scraped_mod.nexus_mod_id)
                .collect::<Vec<i32>>(),
        )
        .await?;
        let mods_to_create_or_update: Vec<UnsavedMod> = scraped
            .mods
            .iter()
            .map(|scraped_mod| UnsavedMod {
                name: scraped_mod.name,
                nexus_mod_id: scraped_mod.nexus_mod_id,
                author_name: scraped_mod.author_name,
                author_id: Some(scraped_mod.author_id),
                category_name: scraped_mod.category_name,
                category_id: scraped_mod.category_id,
                description: scraped_mod.desc,
                thumbnail_link: scraped_mod.thumbnail_link,
                game_id: game.id,
                last_update_at: Some(NaiveDateTime::new(
                    scraped_mod.last_update_at,
                    NaiveTime::from_hms(0, 0, 0),
                )),
                first_upload_at: Some(NaiveDateTime::new(
                    scraped_mod.first_upload_at,
                    NaiveTime::from_hms(0, 0, 0),
                )),
            })
            .collect();

        let mods = game_mod::batched_insert(&pool, &mods_to_create_or_update).await?;

        let mods: Vec<Mod> =
            mods.into_iter()
                .filter(|updated_mod| {
                    if let Some(processed_mod) = processed_mods.iter().find(|processed_mod| {
                        processed_mod.nexus_mod_id == updated_mod.nexus_mod_id
                    }) {
                        if let Some(last_update_at) = updated_mod.last_update_at {
                            if processed_mod.last_updated_files_at > last_update_at {
                                return false;
                            }
                        }
                    }
                    true
                })
                .collect();

        for db_mod in mods {
            let mod_span = info_span!("mod", name = ?&db_mod.name, id = &db_mod.nexus_mod_id);
            let _mod_span = mod_span.enter();
            let files_resp = nexus_api::files::get(&client, db_mod.nexus_mod_id).await?;

            debug!(duration = ?files_resp.wait, "sleeping");
            sleep(files_resp.wait).await;

            // Filter out replaced/deleted files (indicated by null category) and archived files
            let files = files_resp
                .files()?
                .into_iter()
                .filter(|file| match file.category {
                    None => {
                        info!(
                            name = file.file_name,
                            id = file.file_id,
                            "skipping file with no category"
                        );
                        false
                    }
                    Some(category) if category == "ARCHIVED" => false,
                    Some(_) => true,
                });

            let processed_file_ids: HashSet<i32> =
                file::get_processed_nexus_file_ids_by_mod_id(&pool, db_mod.id)
                    .await?
                    .into_iter()
                    .collect();

            for api_file in files {
                let file_span =
                    info_span!("file", name = &api_file.file_name, id = &api_file.file_id,);
                let _file_span = file_span.enter();

                if processed_file_ids.contains(&(api_file.file_id as i32)) {
                    info!("skipping file already present and processed in database");
                    continue;
                }
                let db_file = file::insert(
                    &pool,
                    &file::UnsavedFile {
                        name: api_file.name,
                        file_name: api_file.file_name,
                        nexus_file_id: api_file.file_id as i32,
                        mod_id: db_mod.id,
                        category: api_file.category,
                        version: api_file.version,
                        mod_version: api_file.mod_version,
                        size: api_file.size,
                        uploaded_at: api_file.uploaded_at,
                    },
                )
                .await?;

                let mut checked_metadata = false;
                match nexus_api::metadata::contains_plugin(&client, &api_file).await {
                    Ok(contains_plugin) => {
                        if let Some(contains_plugin) = contains_plugin {
                            checked_metadata = true;
                            if !contains_plugin {
                                info!("file metadata does not contain a plugin, skip downloading");
                                file::update_has_plugin(&pool, db_file.id, false).await?;
                                continue;
                            }
                        } else {
                            warn!("file has no metadata link, continuing with download");
                        }
                    }
                    Err(err) => {
                        warn!(error = %err, "error retreiving metadata for file, continuing with download");
                    }
                };

                let humanized_size = api_file
                    .size
                    .file_size(file_size_opts::CONVENTIONAL)
                    .expect("unable to create human-readable file size");
                info!(size = %humanized_size, "decided to download file");
                let download_link_resp =
                    nexus_api::download_link::get(&client, db_mod.nexus_mod_id, api_file.file_id)
                        .await;
                if let Err(err) = &download_link_resp {
                    if let Some(reqwest_err) = err.downcast_ref::<reqwest::Error>() {
                        if reqwest_err.status() == Some(StatusCode::NOT_FOUND) {
                            warn!(
                                status = ?reqwest_err.status(),
                                "failed to get download link for file, skipping file"
                            );
                            file::update_has_download_link(&pool, db_file.id, false).await?;
                            continue;
                        }
                    }
                }
                let download_link_resp = download_link_resp?;

                let mut tokio_file = match download_link_resp.download_file(&client).await {
                    Ok(file) => {
                        info!(bytes = api_file.size, "download finished");
                        file::update_downloaded_at(&pool, db_file.id).await?;
                        file
                    }
                    Err(err) => {
                        warn!(error = %err, "failed all attempts at downloading file, skipping file");
                        continue;
                    }
                };

                let mut initial_bytes = [0; 8];
                tokio_file.seek(SeekFrom::Start(0)).await?;
                if let Err(err) = tokio_file.read_exact(&mut initial_bytes).await {
                    warn!(error = %err, "failed to read initial bytes, skipping file");
                    file::update_unable_to_extract_plugins(&pool, db_file.id, true).await?;
                    continue;
                }
                let kind = match infer::get(&initial_bytes) {
                    Some(kind) => kind,
                    None => {
                        warn!(initial_bytes = ?initial_bytes, "unable to determine file type of archive, skipping file");
                        file::update_unable_to_extract_plugins(&pool, db_file.id, true).await?;
                        continue;
                    }
                };
                info!(
                    mime_type = kind.mime_type(),
                    "inferred mime_type of downloaded archive"
                );

                match kind.mime_type() {
                    "application/vnd.rar" => {
                        info!("downloaded archive is RAR archive, attempt to uncompress entire archive");
                        // Use unrar to uncompress the entire .rar file to avoid bugs with compress_tools uncompressing certain .rar files:
                        // https://github.com/libarchive/libarchive/issues/373, https://github.com/libarchive/libarchive/issues/1426
                        tokio_file.seek(SeekFrom::Start(0)).await?;
                        let mut file = tokio_file.try_clone().await?.into_std().await;
                        match extract_with_unrar(
                            &mut file,
                            &pool,
                            &db_file,
                            &db_mod,
                            checked_metadata,
                        )
                        .await
                        {
                            Ok(_) => Ok(()),
                            Err(err) => {
                                // unrar failed to extract rar file (e.g. archive has unicode filenames)
                                // Attempt to uncompress the archive using `7z` unix command instead
                                warn!(error = %err, "failed to extract file with unrar, extracting whole archive with 7z instead");
                                extract_with_7zip(&mut file, &pool, &db_file, &db_mod).await
                            }
                        }?;
                    }
                    _ => {
                        tokio_file.seek(SeekFrom::Start(0)).await?;
                        let mut file = tokio_file.try_clone().await?.into_std().await;

                        match extract_with_compress_tools(&mut file, &pool, &db_file, &db_mod).await
                        {
                            Ok(_) => Ok(()),
                            Err(err) => {
                                if err
                                    .downcast_ref::<extractors::compress_tools::ExtractorError>()
                                    .is_some()
                                    && (kind.mime_type() == "application/zip"
                                        || kind.mime_type() == "application/x-7z-compressed")
                                {
                                    // compress_tools or libarchive failed to extract zip/7z file (e.g. archive is deflate64 compressed)
                                    // Attempt to uncompress the archive using `7z` unix command instead
                                    warn!(error = %err, "failed to extract file with compress_tools, extracting whole archive with 7z instead");
                                    extract_with_7zip(&mut file, &pool, &db_file, &db_mod).await
                                } else {
                                    Err(err)
                                }
                            }
                        }?;
                    }
                }

                debug!(duration = ?download_link_resp.wait, "sleeping");
                sleep(download_link_resp.wait).await;
            }

            game_mod::update_last_updated_files_at(&pool, db_mod.id).await?;
        }

        page += 1;
        debug!(?page, ?has_next_page, "sleeping 1 second");
        sleep(Duration::from_secs(1)).await;
    }

    Ok(())
}